add pretrained option for

This commit is contained in:
2022-10-29 20:05:13 +00:00
parent ef05cf57d4
commit 6dd16c7002
2 changed files with 7 additions and 3 deletions
+4
View File
@@ -206,6 +206,10 @@ if __name__ == '__main__':
model : AutoEncoder = AutoEncoder(config['cond_ae']) model : AutoEncoder = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr']) optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
if config['pretrained_model'] != '':
load_model(model.encoder, config['pretrained_model'] + '_enc')
load_model(model.decoder, config['pretrained_model'] + '_dec')
ckpt_enc = tf.train.Checkpoint(model=model.encoder) ckpt_enc = tf.train.Checkpoint(model=model.encoder)
manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20) manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20)
ckpt_dec = tf.train.Checkpoint(model=model.decoder) ckpt_dec = tf.train.Checkpoint(model=model.decoder)
+3 -3
View File
@@ -11,12 +11,12 @@ config['notes'] = 'convnext + convnext, add dense with tanh at the end of revers
' '
config['cond_ae']['upscale_output'] = False config['cond_ae']['upscale_output'] = False
config['batch_size'] = 32 config['pretrained_model'] = ''
config['signature_dynamic'] = { 'id': (), } config['signature_dynamic'] = { 'id': (), }
config['signature_static'] = { 'uv_static': (512, 256, 3), config['signature_static'] = { 'uv_static': (512, 256, 3),
'mesh_static': (14475, 3,), 'mesh_static': (config['smpl_res'], 3,),
'mesh_mask': (14475,), 'mesh_mask': (config['smpl_res'],),
'uv_mask': (512, 256,), 'uv_mask': (512, 256,),
} }