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