From 6dd16c7002634f460809405cd6410423a85911ae Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:05:13 +0000 Subject: [PATCH] add pretrained option for --- auto_encoder_static.py | 4 ++++ configs/config_ae_static.py | 6 +++--- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/auto_encoder_static.py b/auto_encoder_static.py index d1244a2..4c74fb4 100644 --- a/auto_encoder_static.py +++ b/auto_encoder_static.py @@ -206,6 +206,10 @@ if __name__ == '__main__': model : AutoEncoder = AutoEncoder(config['cond_ae']) 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) manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20) ckpt_dec = tf.train.Checkpoint(model=model.decoder) diff --git a/configs/config_ae_static.py b/configs/config_ae_static.py index c653685..317be2f 100644 --- a/configs/config_ae_static.py +++ b/configs/config_ae_static.py @@ -11,12 +11,12 @@ config['notes'] = 'convnext + convnext, add dense with tanh at the end of revers ' config['cond_ae']['upscale_output'] = False -config['batch_size'] = 32 +config['pretrained_model'] = '' config['signature_dynamic'] = { 'id': (), } config['signature_static'] = { 'uv_static': (512, 256, 3), - 'mesh_static': (14475, 3,), - 'mesh_mask': (14475,), + 'mesh_static': (config['smpl_res'], 3,), + 'mesh_mask': (config['smpl_res'],), 'uv_mask': (512, 256,), } \ No newline at end of file