From c29654735cc335f37d81120126708b4e187ba04b Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 30 Jun 2023 18:29:03 +0200 Subject: [PATCH] update training scripts --- configs/ae_body_cape_m.py | 8 ++-- configs/iccvw/ae_body_1.py | 78 +++++++++++++++++++++++++++++++ configs/iccvw/ae_body_cape_f.py | 73 +++++++++++++++++++++++++++++ configs/iccvw/ae_body_cape_m.py | 73 +++++++++++++++++++++++++++++ configs/iccvw/ae_norm_1.py | 78 +++++++++++++++++++++++++++++++ configs/iccvw/ae_static_1.py | 78 +++++++++++++++++++++++++++++++ configs/iccvw/ae_static_cape_f.py | 74 +++++++++++++++++++++++++++++ configs/iccvw/ae_static_cape_m.py | 74 +++++++++++++++++++++++++++++ scripts/ae_body_train.py | 9 ++-- scripts/ae_norm_train.py | 17 +++---- scripts/ae_static_train.py | 11 +++-- 11 files changed, 552 insertions(+), 21 deletions(-) create mode 100644 configs/iccvw/ae_body_1.py create mode 100644 configs/iccvw/ae_body_cape_f.py create mode 100644 configs/iccvw/ae_body_cape_m.py create mode 100644 configs/iccvw/ae_norm_1.py create mode 100644 configs/iccvw/ae_static_1.py create mode 100644 configs/iccvw/ae_static_cape_f.py create mode 100644 configs/iccvw/ae_static_cape_m.py diff --git a/configs/ae_body_cape_m.py b/configs/ae_body_cape_m.py index b7b0a94..a896f40 100644 --- a/configs/ae_body_cape_m.py +++ b/configs/ae_body_cape_m.py @@ -2,7 +2,7 @@ from utils.convnext_utils import CNBlockConfig config = dict() -config['gpu'] = '4' +config['gpu'] = '0' config['name'] = 'body_torch_cape_m' config['notes'] = ' ' config['tags'] = ['torch', 'body'] @@ -11,9 +11,9 @@ config['script'] = 'scripts/ae_body_train.py' config['pretrained_model'] = '' config['batch_size'] = 4 -config['train_dir'] = '/data-local/127/cape_processed/train/male' -config['valid_dir'] = '/data-local/127/cape_processed/test/male' -config['test_dir'] = '/data-local/127/cape_processed/test/male' +config['train_dir'] = '/data-local/128/cape_processed/train/male' +config['valid_dir'] = '/data-local/128/cape_processed/test/male' +config['test_dir'] = '/data-local/128/cape_processed/test/male' config['train_frames'] = None config['valid_frames'] = None diff --git a/configs/iccvw/ae_body_1.py b/configs/iccvw/ae_body_1.py new file mode 100644 index 0000000..a494bff --- /dev/null +++ b/configs/iccvw/ae_body_1.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '2' +config['name'] = 'body_iccvw_1' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_body_cape_f.py b/configs/iccvw/ae_body_cape_f.py new file mode 100644 index 0000000..af816a5 --- /dev/null +++ b/configs/iccvw/ae_body_cape_f.py @@ -0,0 +1,73 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'body_iccvw_cape_f' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/128/cape_processed/train/female' +config['valid_dir'] = '/data-local/128/cape_processed/test/female' +config['test_dir'] = '/data-local/128/cape_processed/test/female' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_body_cape_m.py b/configs/iccvw/ae_body_cape_m.py new file mode 100644 index 0000000..a896f40 --- /dev/null +++ b/configs/iccvw/ae_body_cape_m.py @@ -0,0 +1,73 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '0' +config['name'] = 'body_torch_cape_m' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/128/cape_processed/train/male' +config['valid_dir'] = '/data-local/128/cape_processed/test/male' +config['test_dir'] = '/data-local/128/cape_processed/test/male' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/configs/iccvw/ae_norm_1.py b/configs/iccvw/ae_norm_1.py new file mode 100644 index 0000000..f4cdb16 --- /dev/null +++ b/configs/iccvw/ae_norm_1.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '1' +config['name'] = 'norm_iccvw_1' +config['notes'] = ' ' +config['tags'] = ['torch', 'norm'] +config['script'] = 'scripts/ae_norm_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 40 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_mask', + 'uv_normals', + 'mesh_mask', + 'mesh_normals', + } \ No newline at end of file diff --git a/configs/iccvw/ae_static_1.py b/configs/iccvw/ae_static_1.py new file mode 100644 index 0000000..05836f4 --- /dev/null +++ b/configs/iccvw/ae_static_1.py @@ -0,0 +1,78 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '3' +config['name'] = 'static_iccvw_2' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_train.py' + +config['pretrained_model'] = 'checkpoints/iccvw/static_iccvw_1' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cloth3d_processed' +config['valid_dir'] = '/data-local/127/cloth3d_processed' +config['test_dir'] = '/data-net/dataset/test_ready' + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = False +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +if config['is_skirt']: + config['train_frames'] = 'frames_train_skirt.txt' + config['valid_frames'] = 'frames_valid_skirt.txt' +else: + config['train_frames'] = 'frames_train.txt' + config['valid_frames'] = 'frames_valid.txt' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0005 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/configs/iccvw/ae_static_cape_f.py b/configs/iccvw/ae_static_cape_f.py new file mode 100644 index 0000000..efe2cb2 --- /dev/null +++ b/configs/iccvw/ae_static_cape_f.py @@ -0,0 +1,74 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'static_iccvw_cape_f' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127_2/cape_processed/train/female' +config['valid_dir'] = '/data-local/127_2/cape_processed/test/female' +config['test_dir'] = '/data-local/127_2/cape_processed/test/female' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0005 +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/configs/iccvw/ae_static_cape_m.py b/configs/iccvw/ae_static_cape_m.py new file mode 100644 index 0000000..0685b09 --- /dev/null +++ b/configs/iccvw/ae_static_cape_m.py @@ -0,0 +1,74 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'static_iccvw_cape_m' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127_2/cape_processed/train/male' +config['valid_dir'] = '/data-local/127_2/cape_processed/test/male' +config['test_dir'] = '/data-local/127_2/cape_processed/test/male' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0005 +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/scripts/ae_body_train.py b/scripts/ae_body_train.py index 6226991..e6dfa9a 100644 --- a/scripts/ae_body_train.py +++ b/scripts/ae_body_train.py @@ -14,7 +14,7 @@ args = vars(parser.parse_args()) config_module = importlib.import_module(args['config']) config = config_module.config if args['debug']: - from configs.ae_body_cape_m import config + from configs.iccvw.ae_body_cape_m import config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: @@ -49,6 +49,7 @@ from pytorch_msssim import MS_SSIM def init_wandb(): + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/ae_body_train.py') @@ -198,15 +199,15 @@ if __name__ == '__main__': load_model(model.encoder, config['pretrained_model'] + '_enc') load_model(model.decoder, config['pretrained_model'] + '_dec') - ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5) - ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', 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) run_name = config['name'] print(f'started training: {run_name}') if not args['debug']: # model = torch.compile(model) - wandb.watch(model, log='all', log_freq=1000, log_graph=True) + wandb.watch(model, log='all', log_freq=10000, log_graph=True) train() diff --git a/scripts/ae_norm_train.py b/scripts/ae_norm_train.py index 621b493..5a0b438 100644 --- a/scripts/ae_norm_train.py +++ b/scripts/ae_norm_train.py @@ -49,7 +49,8 @@ from pytorch_msssim import MS_SSIM def init_wandb(): - run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127_2/temp', tags=config['tags']) + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) + run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/ae_norm_train.py') artifact.add_file('losses/losses.py') @@ -70,9 +71,9 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_ def train_step(batch): model.train() - pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1 + pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4 - ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1)) + ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4)) loss = 1e3 * (1 - ms_ssim) optimizer.zero_grad() @@ -93,13 +94,13 @@ def valid_step(batch): model.eval() with torch.inference_mode(): - pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.1 + pred = model([batch['uv_normals'], batch['uv_mask']]) * 1.4 mesh_pred = uv2mesh(pred) mesh_orig = batch['mesh_normals'] l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_normals'] - pred)), axis=1)) - ms_ssim = torch.mean(_MSSSIM(pred + 1.1, batch['uv_normals'] + 1.1)) + ms_ssim = torch.mean(_MSSSIM(pred + 1.4, batch['uv_normals'] + 1.4)) loss = 1e3 * (1 - ms_ssim) metric = torch.mean(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) @@ -109,7 +110,7 @@ def valid_step(batch): def plot_everything(batch: dict, type: str): model.eval() - pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.1 + pred = model([batch['uv_normals'], batch['uv_mask']]).detach() * 1.4 # mesh_pred = uv2mesh(pred) # mesh_orig = batch['mesh_normals'] @@ -117,7 +118,7 @@ def plot_everything(batch: dict, type: str): orig_img = torch.permute(batch['uv_normals'], [0,2,3,1]) pred_img = torch.permute(pred, [0,2,3,1]) img = torch.concat((orig_img, pred_img), axis=2) - imgs = [wandb.Image((img[i].cpu().numpy() + 1.1) / 2.2) for i in range(4)] + imgs = [wandb.Image((img[i].cpu().numpy() + 1.4) / 2.8) for i in range(4)] # F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy") # garment_unposed, faces = mesh_pred.cpu().numpy(), F @@ -196,7 +197,7 @@ if __name__ == '__main__': optimizer = optim.Adam(model.parameters(), lr=config['lr']) # _MSSSIM = MSSSIM(data_range=2).to(config['device']) - _MSSSIM = MS_SSIM(data_range=2.2, size_average=True, channel=3) + _MSSSIM = MS_SSIM(data_range=2.8, size_average=True, channel=3) if config['pretrained_model'] != '': load_model(model.encoder, config['pretrained_model'] + '_enc') diff --git a/scripts/ae_static_train.py b/scripts/ae_static_train.py index 1d03715..6774d52 100644 --- a/scripts/ae_static_train.py +++ b/scripts/ae_static_train.py @@ -13,8 +13,8 @@ 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.ae_static_orig_2 import config + # from configs.ae_static_cape_m import config + from configs.iccvw.ae_static_1 import config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: @@ -49,6 +49,7 @@ from pytorch_msssim import MS_SSIM def init_wandb(): + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True) run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/ae_static_train.py') @@ -208,15 +209,15 @@ if __name__ == '__main__': load_model(model.encoder, config['pretrained_model'] + '_enc') load_model(model.decoder, config['pretrained_model'] + '_dec') - ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5) - ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', 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) run_name = config['name'] print(f'started training: {run_name}') if not args['debug']: # model = torch.compile(model) - wandb.watch(model, log='all', log_freq=1000, log_graph=True) + wandb.watch(model, log='all', log_freq=10000, log_graph=True) train()