updated static training

scripts/ae_static_train.py

@@ -10,11 +10,13 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
-config_module = importlib.import_module(args['config'])
-config = config_module.config
 if args['debug']:
     # from configs.ae_static_cape_m import config
     from configs.iccvw.ae_static_1 import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -206,8 +208,10 @@ if __name__ == '__main__':
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
 
     if config['pretrained_model'] != '':
-        load_model(model.encoder, config['pretrained_model'] + '_enc')
-        load_model(model.decoder, config['pretrained_model'] + '_dec')
+        it = load_model(model.encoder, config['pretrained_model'] + '_enc', optimizer=optimizer)
+        load_model(model.decoder, config['pretrained_model'] + '_dec', optimizer=optimizer)
+        pre_name = config['pretrained_model']
+        print(f'continued training from: {pre_name} at iteration: {it}')
 
     ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
     ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5)