From 64ecb900f5fa2f717c8b7e6089928c24b4cae75a Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 26 Feb 2023 22:25:50 +0000 Subject: [PATCH] started training pyramid --- aaa_helpers.txt | 25 +++- configs/iccv/config_abl_pyr_32.py | 54 +++---- configs/iccv/config_abl_pyr_32_skirt.py | 188 ++++++++++++++++++++++++ cvae.py | 42 ++++-- cvae_pyramid.py | 50 ++++--- models/pyramid_level.py | 7 +- models/pyramid_model_optimized.py | 4 - upscale_test.py | 23 ++- 8 files changed, 314 insertions(+), 79 deletions(-) create mode 100644 configs/iccv/config_abl_pyr_32_skirt.py diff --git a/aaa_helpers.txt b/aaa_helpers.txt index 6e504a1..40535dc 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -14,12 +14,27 @@ nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_m" > nohup/icc nohup python cvae.py --name="lhunor_abl_all" > nohup/iccv/ablation/abl_all.out 2>&1 & -nohup python upscale_test.py --name="iccv_32" > nohup/iccv/upscale/iccv_32.out 2>&1 & -nohup python upscale_test.py --name="iccv_64" > nohup/iccv/upscale/iccv_64.out 2>&1 & -nohup python upscale_test.py --name="iccv_128" > nohup/iccv/upscale/iccv_128.out 2>&1 & -nohup python upscale_test.py --name="iccv_256" > nohup/iccv/upscale/iccv_256.out 2>&1 & +nohup python upscale_test.py --name="iccv_32_2" > nohup/iccv/upscale/iccv_32_2.out 2>&1 & +nohup python upscale_test.py --name="iccv_64_2" > nohup/iccv/upscale/iccv_64_2.out 2>&1 & +nohup python upscale_test.py --name="iccv_128_2" > nohup/iccv/upscale/iccv_128_2.out 2>&1 & +nohup python upscale_test.py --name="iccv_256_2" > nohup/iccv/upscale/iccv_256_2.out 2>&1 & -nohup python upscale_test.py --name="iccv_32_skirt" > nohup/iccv/upscale/iccv_32_skirt.out 2>&1 & +nohup python upscale_test.py --name="iccv_32_2_skirt" > nohup/iccv/upscale/iccv_32_2_skirt.out 2>&1 & +nohup python upscale_test.py --name="iccv_64_2_skirt" > nohup/iccv/upscale/iccv_64_2_skirt.out 2>&1 & +nohup python upscale_test.py --name="iccv_128_2_skirt" > nohup/iccv/upscale/iccv_128_2_skirt.out 2>&1 & +nohup python upscale_test.py --name="iccv_256_2_skirt" > nohup/iccv/upscale/iccv_256_2_skirt.out 2>&1 & + +nohup python upscale_test.py --name="iccv_32_2_cape_f" > nohup/iccv/upscale/iccv_32_2_cape_f_real.out 2>&1 & +nohup python upscale_test.py --name="iccv_64_2_cape_f" > nohup/iccv/upscale/iccv_64_2_cape_f.out 2>&1 & +nohup python upscale_test.py --name="iccv_128_2_cape_f" > nohup/iccv/upscale/iccv_128_2_cape_f.out 2>&1 & +nohup python upscale_test.py --name="iccv_256_2_cape_f" > nohup/iccv/upscale/iccv_256_2_cape_f.out 2>&1 & + +nohup python upscale_test.py --name="iccv_32_2_cape_m" > nohup/iccv/upscale/iccv_32_2_cape_m.out 2>&1 & +nohup python upscale_test.py --name="iccv_64_2_cape_m" > nohup/iccv/upscale/iccv_64_2_cape_m.out 2>&1 & +nohup python upscale_test.py --name="iccv_128_2_cape_m" > nohup/iccv/upscale/iccv_128_2_cape_m.out 2>&1 & +nohup python upscale_test.py --name="iccv_256_2_cape_m" > nohup/iccv/upscale/iccv_256_2_cape_m.out 2>&1 & + +nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32" > nohup/iccv/pyr/iccv_pyr_32.out 2>&1 & skirts!!!! pyramid next week diff --git a/configs/iccv/config_abl_pyr_32.py b/configs/iccv/config_abl_pyr_32.py index 3408e7f..a8d4070 100644 --- a/configs/iccv/config_abl_pyr_32.py +++ b/configs/iccv/config_abl_pyr_32.py @@ -1,6 +1,6 @@ config = dict() -config['gpu'] = '2' +config['gpu'] = '0' config['name'] = 'iccv_pyr_32' config['notes'] = '' config['tags'] = ['pyr', 'iccv'] @@ -47,6 +47,7 @@ config['lr'] = 0.0001 config['starting_balancing_coeff'] = 0.001 config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 config['signature_dynamic'] = { 'id': (), 'frame_nr': (), 'pose_shape': (72 + 11,), @@ -63,8 +64,8 @@ config['signature_static'] = { 'uv_static': (512, 256, 3), } config['pyramid']['is_label'] = False -config['pyramid']['add_noise'] = True -config['pyramid']['noise_amount'] = 0.001 +config['pyramid']['add_noise'] = False +config['pyramid']['noise_amount'] = 0.0 config['pyramid']['loss_weights'] = [1,1,1,1,1] @@ -77,9 +78,10 @@ for res in config['pyramid']['resolutions']: config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3) config['signature_static']['uv_mask' + res_modifier] = (res, res//2) -config['pyramid']['static_model'] = 'checkpoints/post/static_post_2' -config['pyramid']['body_model'] = 'checkpoints/post/body_post_2' -config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1' + +config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1' +config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1' +config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1' config['pyramid']['freeze_encoders'] = True @@ -88,10 +90,10 @@ for res in config['pyramid']['resolutions']: if res == 512: level = dict() level['resolution'] = 512 - level['latent_dim'] = config['latent_dim'] + level['latent_dim'] = 64 * config['bottleneck_scale'] level['kernel_size'] = 7 level['encoder_depths'] = [3, 3, 9, 3] - level['encoder_dims'] = [32, 64, 128, 64] + level['encoder_dims'] = [32, 64, 128, level['latent_dim']] level['decoder_input'] = (16, 8, 80) level['decoder_depths'] = [3, 3, 9, 3] level['decoder_dims'] = [256, 128, 64, 16] @@ -101,53 +103,55 @@ for res in config['pyramid']['resolutions']: elif res == 256: level = dict() level['resolution'] = 256 - level['latent_dim'] = config['latent_dim'] - level['kernel_size'] = 7 + level['latent_dim'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 5 level['encoder_depths'] = [9, 9, 9, 9] - level['encoder_dims'] = [32, 64, 128, 64] + level['encoder_dims'] = [32, 64, 128, level['latent_dim']] level['decoder_input'] = (8, 4, 128) level['decoder_depths'] = [9, 9, 9, 9] level['decoder_dims'] = [256, 128, 64, 16] level['avg_garment_path'] = '' - level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2' + level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#'checkpoints/up_3_d256_512_2' level['upscale_enhancer_freeze'] = True config['pyramid']['levels'].append(level) elif res == 128: level = dict() level['resolution'] = 128 - level['latent_dim'] = config['latent_dim'] + level['latent_dim'] = 64 * config['bottleneck_scale'] level['kernel_size'] = 5 level['encoder_depths'] = [16, 9, 9, 9] level['encoder_dims'] = [64, 128, 128, 64] level['decoder_input'] = (4, 2, 128) level['decoder_depths'] = [16, 9, 9, 9] - level['decoder_dims'] = [256, 128, 128, 64] + level['decoder_dims'] = [256, 128, 128, level['latent_dim']] level['avg_garment_path'] = '' - level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2' + level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#'checkpoints/up_3_d128_256_2' level['upscale_enhancer_freeze'] = True config['pyramid']['levels'].append(level) elif res == 64: level = dict() level['resolution'] = 64 - level['kernel_size'] = 5 + level['latent_dim'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 3 level['encoder_depths'] = [9, 9] - level['encoder_dims'] = [64, 64] + level['encoder_dims'] = [64, level['latent_dim']] level['decoder_depths'] = level['encoder_depths'] level['decoder_dims'] = level['encoder_dims'] level['avg_garment_path'] = '' - level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2' + level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#checkpoints/up_3_d64_128_2' level['upscale_enhancer_freeze'] = True config['pyramid']['levels'].append(level) elif res == 32: level = dict() level['resolution'] = 32 + level['latent_dim'] = 64 * config['bottleneck_scale'] level['kernel_size'] = 3 level['encoder_depths'] = [16] - level['encoder_dims'] = [64] + level['encoder_dims'] = [level['latent_dim']] level['decoder_depths'] = level['encoder_depths'] level['decoder_dims'] = level['encoder_dims'] level['avg_garment_path'] = '' - level['upscale_enhancer'] = 'checkpoints/32_512_4'#'checkpoints/up_3_d32_64_2' + level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2' level['upscale_enhancer_freeze'] = True config['pyramid']['levels'].append(level) else: @@ -156,9 +160,9 @@ for res in config['pyramid']['resolutions']: config['cond_ae_static'] = dict() config['cond_ae_static']['is_label'] = False config['cond_ae_static']['kernel_size'] = 7 -config['cond_ae_static']['bottleneck_size'] = 8 -config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3] -config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']] +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] +config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3] +config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']] config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths'] config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims'] config['cond_ae_static']['activation'] = 'tanh' @@ -166,7 +170,7 @@ config['cond_ae_static']['activation'] = 'tanh' config['cond_ae_curv'] = dict() config['cond_ae_curv']['is_label'] = True config['cond_ae_curv']['kernel_size'] = 7 -config['cond_ae_curv']['bottleneck_size'] = 40 +config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale'] config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3] config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']] config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths'] @@ -176,7 +180,7 @@ config['cond_ae_curv']['activation'] = 'tanh' config['cond_ae_body'] = dict() config['cond_ae_body']['is_label'] = False config['cond_ae_body']['kernel_size'] = 7 -config['cond_ae_body']['bottleneck_size'] = 16 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3] config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']] config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths'] diff --git a/configs/iccv/config_abl_pyr_32_skirt.py b/configs/iccv/config_abl_pyr_32_skirt.py new file mode 100644 index 0000000..881cd2c --- /dev/null +++ b/configs/iccv/config_abl_pyr_32_skirt.py @@ -0,0 +1,188 @@ +config = dict() + +config['gpu'] = '0' +config['name'] = 'iccv_pyr_32_skirt' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv'] + +# DON'T FORGET TO RESET SSIM LOSS USAGE +# DON'T FORGET TO CHECK MODEL LOADING +# DON'T FORGET TO CHECK IF STORE IN MEMORY +# DON'T FORGET TO CHECK REG LOSS COEFF +# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512 +# DON'T FORGET TO CHECK IF SKIRT +# DON'T FORGET TO CHECK IF CAPE +# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2 +# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS + +config['use_skinning'] = True # and deshaping +config['remove_global_orientation'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_offsets'] = False +config['use_norm_loss'] = True + +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +config['pyramid'] = dict() +config['pyramid']['resolutions'] = [32] +config['pyramid']['pretrained_pyramid'] = '' +config['pyramid']['final_upscale_enhancer'] = '' + +config['train_dir'] = '/home/workspace/data/dataset/local-cache/train_skirt' +config['valid_dir'] = '/home/workspace/local-data/valid_skirt' +config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2' +config['batch_size'] = 4 +config['is_skirt'] = False +config['is_cape'] = False + +config['dataset_in_memory'] = False +config['resolution'] = None + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['lr'] = 0.0001 +config['starting_balancing_coeff'] = 0.001 +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision +config['avg_garment_path'] = '' +config['bottleneck_scale'] = 2 + +config['signature_dynamic'] = { 'id': (), 'frame_nr': (), + 'pose_shape': (72 + 11,), + 'uv_body_posed': (512, 256, 3), + 'mesh_body_posed': (config['smpl_res'], 3,), + 'mesh_posed': (config['smpl_res'], 3,), + 'mesh_unposed': (config['smpl_res'], 3,), + 'uv_unposed': (512, 256, 3), + 'uv_curvature': (512, 256, 7), + } +config['signature_static'] = { 'uv_static': (512, 256, 3), + 'mesh_mask': (config['smpl_res'],), + 'uv_mask': (512, 256,), + } + +config['pyramid']['is_label'] = False +config['pyramid']['add_noise'] = False +config['pyramid']['noise_amount'] = 0.0 + +config['pyramid']['loss_weights'] = [1,1,1,1,1] + +for res in config['pyramid']['resolutions']: + if res != 512: + res_modifier = '_' + str(res) + config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3) + config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3) + config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7) + config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3) + config['signature_static']['uv_mask' + res_modifier] = (res, res//2) + + +config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1' +config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1' +config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1' +config['pyramid']['freeze_encoders'] = True + + +config['pyramid']['levels'] = [] +for res in config['pyramid']['resolutions']: + if res == 512: + level = dict() + level['resolution'] = 512 + level['latent_dim'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 7 + level['encoder_depths'] = [3, 3, 9, 3] + level['encoder_dims'] = [32, 64, 128, level['latent_dim']] + level['decoder_input'] = (16, 8, 80) + level['decoder_depths'] = [3, 3, 9, 3] + level['decoder_dims'] = [256, 128, 64, 16] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = '' + config['pyramid']['levels'].append(level) + elif res == 256: + level = dict() + level['resolution'] = 256 + level['latent_dim'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 5 + level['encoder_depths'] = [9, 9, 9, 9] + level['encoder_dims'] = [32, 64, 128, level['latent_dim']] + level['decoder_input'] = (8, 4, 128) + level['decoder_depths'] = [9, 9, 9, 9] + level['decoder_dims'] = [256, 128, 64, 16] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#'checkpoints/up_3_d256_512_2' + level['upscale_enhancer_freeze'] = True + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['latent_dim'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 5 + level['encoder_depths'] = [16, 9, 9, 9] + level['encoder_dims'] = [64, 128, 128, 64] + level['decoder_input'] = (4, 2, 128) + level['decoder_depths'] = [16, 9, 9, 9] + level['decoder_dims'] = [256, 128, 128, level['latent_dim']] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#'checkpoints/up_3_d128_256_2' + level['upscale_enhancer_freeze'] = True + config['pyramid']['levels'].append(level) + elif res == 64: + level = dict() + level['resolution'] = 64 + level['latent_dim'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 3 + level['encoder_depths'] = [9, 9] + level['encoder_dims'] = [64, level['latent_dim']] + level['decoder_depths'] = level['encoder_depths'] + level['decoder_dims'] = level['encoder_dims'] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#checkpoints/up_3_d64_128_2' + level['upscale_enhancer_freeze'] = True + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['latent_dim'] = 64 * config['bottleneck_scale'] + level['kernel_size'] = 3 + level['encoder_depths'] = [16] + level['encoder_dims'] = [level['latent_dim']] + level['decoder_depths'] = level['encoder_depths'] + level['decoder_dims'] = level['encoder_dims'] + level['avg_garment_path'] = '' + level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2' + level['upscale_enhancer_freeze'] = True + config['pyramid']['levels'].append(level) + else: + raise NotImplementedError + +config['cond_ae_static'] = dict() +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale'] +config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3] +config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']] +config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths'] +config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims'] +config['cond_ae_static']['activation'] = 'tanh' + +config['cond_ae_curv'] = dict() +config['cond_ae_curv']['is_label'] = True +config['cond_ae_curv']['kernel_size'] = 7 +config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale'] +config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3] +config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']] +config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths'] +config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims'] +config['cond_ae_curv']['activation'] = 'tanh' + +config['cond_ae_body'] = dict() +config['cond_ae_body']['is_label'] = False +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] +config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3] +config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']] +config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths'] +config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims'] +config['cond_ae_body']['activation'] = 'tanh' \ No newline at end of file diff --git a/cvae.py b/cvae.py index 8471adb..7964455 100644 --- a/cvae.py +++ b/cvae.py @@ -42,6 +42,7 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mes from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL from models.cvae_model import CVAENew as CVAE +from custom_utils.custom_utils import Accumulator def init_wandb(): @@ -529,7 +530,8 @@ if __name__ == '__main__': ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + config['name'], max_to_keep=20) - + train_accumulator = Accumulator() + valid_accumulator = Accumulator() if config['pretrained_model'] != '': load_model(model, config['pretrained_model']) @@ -542,31 +544,45 @@ if __name__ == '__main__': batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) train_loss = train_step(batch) + train_accumulator.update(train_loss) - if step % 10 == 0: + if step % 50 == 0: data_valid = valid_iterator.get_next() data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) valid_loss = valid_step(data_valid) + valid_accumulator.update(valid_loss) - to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } - if step % 5000 == 0 or (step < 5000 and step % 500 == 0): - train_figs = plot_everything(batch, 'train') - valid_figs = plot_everything(data_valid, 'valid') - to_log = { **to_log, **train_figs, **valid_figs } + # to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } + # if step % 5000 == 0 or (step < 5000 and step % 500 == 0): + # train_figs = plot_everything(batch, 'train') + # valid_figs = plot_everything(data_valid, 'valid') + # to_log = { **to_log, **train_figs, **valid_figs } - if step % 5000 == 0: - manager.save(checkpoint_number=step) - if step == 0: - to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } + # if step % 5000 == 0: + # manager.save(checkpoint_number=step) + # if step == 0: + # to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017): if balancing_activated_step == 0: balancing_activated_step = step g_balancing_coeff = g_balancing_coeff + 1 / 10000 - to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } } - wandb.log(to_log, step = step) + if step % 1000 == 0 or (step < 1000 and step % 100 == 0): + to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } }) + if step % 5000 == 0 or (step < 5000 and step % 1000 == 0): + train_figs = plot_everything(batch, 'train') + valid_figs = plot_everything(data_valid, 'valid') + to_log = { **to_log, **train_figs, **valid_figs } + if step == 0: + to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } + + to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } } + + wandb.log(to_log, step=step) + train_accumulator.reset() + valid_accumulator.reset() print(step) diff --git a/cvae_pyramid.py b/cvae_pyramid.py index cc4cd55..4010495 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -1,6 +1,6 @@ import os import sys -from configs.post.config_abl_pyr_64_128_5 import config +from configs.iccv.config_abl_pyr_32 import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -43,10 +43,11 @@ from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL from models.pyramid_model_optimized import PyramidNew as Pyramid from postprocessing import fix_collision +from custom_utils.custom_utils import Accumulator def init_wandb(): - run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags']) + run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae_pyramid.py') @@ -61,11 +62,14 @@ def init_wandb(): artifact.add_file('convnext/models/convnext_tf.py') config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.post.')] for config_file_name in config_file_names: - artifact.add_file('/'.join(config_file_name.split('.')) + '.py') + file_name = '/'.join(config_file_name.split('.')) + '.py' + if os.path.exists(file_name): + artifact.add_file(file_name) run.log_artifact(artifact) uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color): fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ]) @@ -791,6 +795,8 @@ if __name__ == '__main__': ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() # load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000') @@ -846,31 +852,33 @@ if __name__ == '__main__': train_loss = train_step(batch) + train_accumulator.update(train_loss) - data_valid = valid_iterator.get_next() - data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) - valid_loss = valid_step(data_valid) + if step % 50 == 0: + data_valid = valid_iterator.get_next() + data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) + valid_loss = valid_step(data_valid) + valid_accumulator.update(valid_loss) - to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } - if step % 5000 == 0 or (step < 5000 and step % 500 == 0): - train_figs = plot_everything(batch, 'train') - valid_figs = plot_everything(data_valid, 'valid') - to_log = { **to_log, **train_figs, **valid_figs } - - if step % 5000 == 0: - manager.save(checkpoint_number=step) - if step == 0: - to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } - - if train_loss['train/euclid_error'] < 0.004: + if train_loss['train/euclid_error'] < 0.015: if balancing_activated_step == 0: balancing_activated_step = step balancing_coeff = balancing_coeff + 1 / 20000 + if step % 1000 == 0 or (step < 1000 and step % 100 == 0): + to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } }) + if step % 5000 == 0 or (step < 5000 and step % 1000 == 0): + train_figs = plot_everything(batch, 'train') + valid_figs = plot_everything(data_valid, 'valid') + to_log = { **to_log, **train_figs, **valid_figs } + if step == 0: + to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } - to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } - - wandb.log(to_log, step = step) + to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } + + wandb.log(to_log, step=step) + train_accumulator.reset() + valid_accumulator.reset() print(step) diff --git a/models/pyramid_level.py b/models/pyramid_level.py index 9cf7322..2d94698 100644 --- a/models/pyramid_level.py +++ b/models/pyramid_level.py @@ -102,13 +102,12 @@ class PyramidLevelNew(tf.keras.Model): super(PyramidLevelNew, self).__init__() self.size = config['resolution'] - self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size']) + self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'], latent_dim=config['latent_dim']) self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size']) self.upscale_enhancer = None if config['upscale_enhancer'] != '': - k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5)) - self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size) + self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=config['kernel_size']) load_model(self.upscale_enhancer, config['upscale_enhancer']) self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze'] @@ -117,7 +116,7 @@ class PyramidLevelNew(tf.keras.Model): else: self.avg_garment = None - self.reshape = tf.keras.layers.Reshape((8, 4, 64)) + self.reshape = tf.keras.layers.Reshape((8, 4, config['latent_dim'])) def enhance(self, low, mask, training=None): diff --git a/models/pyramid_model_optimized.py b/models/pyramid_model_optimized.py index 5301872..0e04ab7 100644 --- a/models/pyramid_model_optimized.py +++ b/models/pyramid_model_optimized.py @@ -343,11 +343,7 @@ class PyramidNew(tf.keras.Model): for i in range(1, len(self.levels)): input = levels_input[i] input[0] = input[0] - pred - input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3]) - input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None] - input[0] = input[0] * input[-2][:,:,:,None] # uv_mask_res pred_level, mean, logvar = self.levels[i](input, z_cond, training=training) - pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res pred = pred + pred_level if i != len(self.levels) - 1: pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training) diff --git a/upscale_test.py b/upscale_test.py index 2900c20..172e7ec 100644 --- a/upscale_test.py +++ b/upscale_test.py @@ -36,14 +36,17 @@ from models.upscaler import UpscalerBasicConv as Upscaler # !!!!!!!!!! CHANGE SKIRT OR NOT -low_res = 128 +low_res = 256 is_skirt = True is_cape = False +is_male = False skirt_modifier = '_skirt' if is_skirt else '' +cape_modifier = '_cape' if is_cape else '' +gender_modifier = '' if not is_cape else ('_m' if is_male else '_f') smpl_res = 3638 if is_cape else 14475 -name=f'iccv_{low_res}_2' + skirt_modifier +name=f'iccv_{low_res}_2' + skirt_modifier + cape_modifier + gender_modifier def init_wandb(): run = wandb.init(project='upscale', name=name, notes='\ @@ -202,7 +205,6 @@ def train(): if step % 5000 == 0: manager.save(checkpoint_number=step) - wandb.log(to_log, step = step) print(step) @@ -255,10 +257,17 @@ if __name__ == '__main__': static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/valid_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) if is_cape: - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) - dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) - static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + if is_male: + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) + static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + else: + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/female', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/female', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/female', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) + static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/female', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + if low_res == 32 or low_res == 64: model = Upscaler(low_res=low_res, kernel_size=3)