add pretrained option for

2022-10-29 20:05:13 +00:00
parent ef05cf57d4
commit 6dd16c7002
2 changed files with 7 additions and 3 deletions
@@ -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)