From cac04a1610f4d5c0735338076b7454b40a0f81da Mon Sep 17 00:00:00 2001
From: Hunor Laczko <luni741@gmail.com>
Date: Sat, 1 Jul 2023 10:46:43 +0200
Subject: [PATCH] updated static training

---
 configs/iccvw/ae_static_1.py      |  6 +++---
 configs/iccvw/ae_static_cape_f.py |  3 +--
 configs/iccvw/ae_static_cape_m.py |  3 +--
 scripts/ae_static_train.py        | 12 ++++++++----
 4 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/configs/iccvw/ae_static_1.py b/configs/iccvw/ae_static_1.py
index 05836f4..a80de57 100644
--- a/configs/iccvw/ae_static_1.py
+++ b/configs/iccvw/ae_static_1.py
@@ -3,12 +3,12 @@ from utils.convnext_utils import CNBlockConfig
 config = dict()
 
 config['gpu'] = '3'
-config['name'] = 'static_iccvw_2'
+config['name'] = 'static_iccvw_4'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'static'] 
 config['script'] = 'scripts/ae_static_train.py'
 
-config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_1'
+config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_3'
 
 config['batch_size'] = 4
 config['train_dir'] = '/data-local/127/cloth3d_processed'
@@ -17,7 +17,7 @@ config['test_dir'] = '/data-net/dataset/test_ready'
 
 config['resolution'] = 512
 config['is_skirt'] = False
-config['dataset_in_memory'] = True
+config['dataset_in_memory'] = False
 config['is_cape'] = False
 config['remove_global_orientation'] = True
 config['device'] = 'cuda'
diff --git a/configs/iccvw/ae_static_cape_f.py b/configs/iccvw/ae_static_cape_f.py
index efe2cb2..c7a928d 100644
--- a/configs/iccvw/ae_static_cape_f.py
+++ b/configs/iccvw/ae_static_cape_f.py
@@ -19,7 +19,7 @@ config['valid_frames'] = None
 
 config['resolution'] = 512
 config['is_skirt'] = False
-config['dataset_in_memory'] = True
+config['dataset_in_memory'] = False
 config['is_cape'] = True
 config['remove_global_orientation'] = True
 config['device'] = 'cuda'
@@ -63,7 +63,6 @@ config['cond_ae']['activation'] = 'tanh'
 
 config['lr'] = 0.0005
 config['cond_ae']['upscale_output'] = False
-config['pretrained_model'] = ''
 
 config['signature'] = { 'id',
                         'frame_nr',
diff --git a/configs/iccvw/ae_static_cape_m.py b/configs/iccvw/ae_static_cape_m.py
index 0685b09..25e2662 100644
--- a/configs/iccvw/ae_static_cape_m.py
+++ b/configs/iccvw/ae_static_cape_m.py
@@ -19,7 +19,7 @@ config['valid_frames'] = None
 
 config['resolution'] = 512
 config['is_skirt'] = False
-config['dataset_in_memory'] = True
+config['dataset_in_memory'] = False
 config['is_cape'] = True
 config['remove_global_orientation'] = True
 config['device'] = 'cuda'
@@ -63,7 +63,6 @@ config['cond_ae']['activation'] = 'tanh'
 
 config['lr'] = 0.0005
 config['cond_ae']['upscale_output'] = False
-config['pretrained_model'] = ''
 
 config['signature'] = { 'id',
                         'frame_nr',
diff --git a/scripts/ae_static_train.py b/scripts/ae_static_train.py
index 6774d52..d445d40 100644
--- a/scripts/ae_static_train.py
+++ b/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)