updated static training

This commit is contained in:
2023-07-01 10:46:43 +02:00
parent f9ac7fe3c3
commit cac04a1610
4 changed files with 13 additions and 11 deletions
+3 -3
View File
@@ -3,12 +3,12 @@ from utils.convnext_utils import CNBlockConfig
config = dict() config = dict()
config['gpu'] = '3' config['gpu'] = '3'
config['name'] = 'static_iccvw_2' config['name'] = 'static_iccvw_4'
config['notes'] = ' ' config['notes'] = ' '
config['tags'] = ['torch', 'static'] config['tags'] = ['torch', 'static']
config['script'] = 'scripts/ae_static_train.py' 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['batch_size'] = 4
config['train_dir'] = '/data-local/127/cloth3d_processed' config['train_dir'] = '/data-local/127/cloth3d_processed'
@@ -17,7 +17,7 @@ config['test_dir'] = '/data-net/dataset/test_ready'
config['resolution'] = 512 config['resolution'] = 512
config['is_skirt'] = False config['is_skirt'] = False
config['dataset_in_memory'] = True config['dataset_in_memory'] = False
config['is_cape'] = False config['is_cape'] = False
config['remove_global_orientation'] = True config['remove_global_orientation'] = True
config['device'] = 'cuda' config['device'] = 'cuda'
+1 -2
View File
@@ -19,7 +19,7 @@ config['valid_frames'] = None
config['resolution'] = 512 config['resolution'] = 512
config['is_skirt'] = False config['is_skirt'] = False
config['dataset_in_memory'] = True config['dataset_in_memory'] = False
config['is_cape'] = True config['is_cape'] = True
config['remove_global_orientation'] = True config['remove_global_orientation'] = True
config['device'] = 'cuda' config['device'] = 'cuda'
@@ -63,7 +63,6 @@ config['cond_ae']['activation'] = 'tanh'
config['lr'] = 0.0005 config['lr'] = 0.0005
config['cond_ae']['upscale_output'] = False config['cond_ae']['upscale_output'] = False
config['pretrained_model'] = ''
config['signature'] = { 'id', config['signature'] = { 'id',
'frame_nr', 'frame_nr',
+1 -2
View File
@@ -19,7 +19,7 @@ config['valid_frames'] = None
config['resolution'] = 512 config['resolution'] = 512
config['is_skirt'] = False config['is_skirt'] = False
config['dataset_in_memory'] = True config['dataset_in_memory'] = False
config['is_cape'] = True config['is_cape'] = True
config['remove_global_orientation'] = True config['remove_global_orientation'] = True
config['device'] = 'cuda' config['device'] = 'cuda'
@@ -63,7 +63,6 @@ config['cond_ae']['activation'] = 'tanh'
config['lr'] = 0.0005 config['lr'] = 0.0005
config['cond_ae']['upscale_output'] = False config['cond_ae']['upscale_output'] = False
config['pretrained_model'] = ''
config['signature'] = { 'id', config['signature'] = { 'id',
'frame_nr', 'frame_nr',
+8 -4
View File
@@ -10,11 +10,13 @@ parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true') parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args()) args = vars(parser.parse_args())
config_module = importlib.import_module(args['config'])
config = config_module.config
if args['debug']: if args['debug']:
# from configs.ae_static_cape_m import config # from configs.ae_static_cape_m import config
from configs.iccvw.ae_static_1 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_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']: if args['debug']:
@@ -206,8 +208,10 @@ if __name__ == '__main__':
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
if config['pretrained_model'] != '': if config['pretrained_model'] != '':
load_model(model.encoder, config['pretrained_model'] + '_enc') it = load_model(model.encoder, config['pretrained_model'] + '_enc', optimizer=optimizer)
load_model(model.decoder, config['pretrained_model'] + '_dec') 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_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5) ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_dec', top_n=5)