From aa96c6eb9cc5d08054bc2728536b08e869d249a8 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 25 Feb 2023 22:46:15 +0000 Subject: [PATCH] more training --- aaa_helpers.txt | 9 +- auto_encoder_static.py | 4 +- configs/iccv/config_abl_pyr_32.py | 184 ++++++++++++++++++++++++ configs/iccv/config_ae_static_cape_m.py | 4 +- models/upscaler.py | 8 +- upscale_test.py | 57 +++++--- 6 files changed, 234 insertions(+), 32 deletions(-) create mode 100644 configs/iccv/config_abl_pyr_32.py diff --git a/aaa_helpers.txt b/aaa_helpers.txt index d6d651a..6e504a1 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -7,13 +7,20 @@ nohup python auto_encoder_static.py --name="lhunor_static_iccv_1_cape_m" > nohup nohup python auto_encoder_body.py --name="lhunor_body_iccv_1" > nohup/iccv/encoders/body_iccv_1.out 2>&1 & nohup python auto_encoder_body.py --name="lhunor_body_iccv_1_skirt" > nohup/iccv/encoders/body_iccv_1_skirt.out 2>&1 & nohup python auto_encoder_body.py --name="lhunor_body_iccv_1_cape_m" > nohup/iccv/encoders/body_iccv_1_cape_m.out 2>&1 & -nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1 " > nohup/iccv/encoders/curv_iccv_1.out 2>&1 & +nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1" > nohup/iccv/encoders/curv_iccv_1.out 2>&1 & nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_skirt" > nohup/iccv/encoders/curv_iccv_1_skirt.out 2>&1 & nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_f" > nohup/iccv/encoders/curv_iccv_1_cape_f.out 2>&1 & nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_m" > nohup/iccv/encoders/curv_iccv_1_cape_m.out 2>&1 & 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_skirt" > nohup/iccv/upscale/iccv_32_skirt.out 2>&1 & + skirts!!!! pyramid next week diff --git a/auto_encoder_static.py b/auto_encoder_static.py index ea709c4..4008917 100644 --- a/auto_encoder_static.py +++ b/auto_encoder_static.py @@ -247,9 +247,9 @@ if __name__ == '__main__': load_model(model.decoder, config['pretrained_model'] + '_dec') ckpt_enc = tf.train.Checkpoint(model=model.encoder) - manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/post/' + config['name'] + '_enc', max_to_keep=20) + manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/iccv/' + config['name'] + '_enc', max_to_keep=20) ckpt_dec = tf.train.Checkpoint(model=model.decoder) - manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/post/' + config['name'] + '_dec', max_to_keep=20) + manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/iccv/' + config['name'] + '_dec', max_to_keep=20) # test() train() \ No newline at end of file diff --git a/configs/iccv/config_abl_pyr_32.py b/configs/iccv/config_abl_pyr_32.py new file mode 100644 index 0000000..3408e7f --- /dev/null +++ b/configs/iccv/config_abl_pyr_32.py @@ -0,0 +1,184 @@ +config = dict() + +config['gpu'] = '2' +config['name'] = 'iccv_pyr_32' +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/local-data/train' +config['valid_dir'] = '/home/workspace/local-data/valid' +config['test_dir'] = '/home/workspace/data/dataset/test_ready' +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['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'] = True +config['pyramid']['noise_amount'] = 0.001 + +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/post/static_post_2' +config['pyramid']['body_model'] = 'checkpoints/post/body_post_2' +config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_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'] = config['latent_dim'] + level['kernel_size'] = 7 + level['encoder_depths'] = [3, 3, 9, 3] + level['encoder_dims'] = [32, 64, 128, 64] + 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'] = config['latent_dim'] + level['kernel_size'] = 7 + level['encoder_depths'] = [9, 9, 9, 9] + level['encoder_dims'] = [32, 64, 128, 64] + 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_freeze'] = True + config['pyramid']['levels'].append(level) + elif res == 128: + level = dict() + level['resolution'] = 128 + level['latent_dim'] = config['latent_dim'] + 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['avg_garment_path'] = '' + level['upscale_enhancer'] = ''#'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['encoder_depths'] = [9, 9] + level['encoder_dims'] = [64, 64] + 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_freeze'] = True + config['pyramid']['levels'].append(level) + elif res == 32: + level = dict() + level['resolution'] = 32 + level['kernel_size'] = 3 + level['encoder_depths'] = [16] + level['encoder_dims'] = [64] + 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_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['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']['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['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['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/configs/iccv/config_ae_static_cape_m.py b/configs/iccv/config_ae_static_cape_m.py index c798b22..de5aff1 100644 --- a/configs/iccv/config_ae_static_cape_m.py +++ b/configs/iccv/config_ae_static_cape_m.py @@ -33,14 +33,14 @@ config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims'] config['cond_ae']['activation'] = 'tanh' -config['gpu'] = '5' +config['gpu'] = '3' config['lr'] = 0.0005 config['name'] = 'static_iccv_1_cape_m' config['notes'] = ' ' config['tags'] = ['iccv', 'static', 'cape_m'] config['cond_ae']['upscale_output'] = False -config['pretrained_model'] = '' +config['pretrained_model'] = 'checkpoints/iccv/static_iccv_1_cape_m_old' config['signature_dynamic'] = { 'id': (), 'frame_nr': (), } diff --git a/models/upscaler.py b/models/upscaler.py index c03e663..187696f 100644 --- a/models/upscaler.py +++ b/models/upscaler.py @@ -131,13 +131,13 @@ class UpscalerBasicConv(tf.keras.Model): self.low_res = low_res self.seq = tf.keras.Sequential([ - tf.keras.layers.Conv2D(256, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'), + tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'), tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'), - tf.keras.layers.Conv2D(128, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'), + tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'), tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'), - tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'), + tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'), tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'), - tf.keras.layers.Conv2D(32, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'), + tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'), tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'), tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'), ]) diff --git a/upscale_test.py b/upscale_test.py index 2e5b6e1..2e4d158 100644 --- a/upscale_test.py +++ b/upscale_test.py @@ -1,5 +1,5 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "6" +os.environ["CUDA_VISIBLE_DEVICES"] = "2" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') @@ -35,7 +35,15 @@ from models.upscaler import UpscalerBasicConv as Upscaler # !!!!!!!!!! CHANGE SSIM BASED ON RES # !!!!!!!!!! CHANGE SKIRT OR NOT -name='256_512_skirt' + +low_res = 64 +is_skirt = True +is_cape = False +skirt_modifier = '_skirt' if is_skirt else '' +smpl_res = 3638 if is_cape else 14475 + + +name=f'iccv_{low_res}_skirt' def init_wandb(): run = wandb.init(project='upscale', name=name, notes='\ @@ -59,15 +67,10 @@ smaller features | \ artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) -is_skirt = True -is_cape = False -skirt_modifier = '_skirt' if is_skirt else '' -smpl_res = 3638 if is_cape else 14475 uv2mesh = UV2Mesh(res=256, is_skirt=is_skirt, is_smpl_orig=is_cape) - def plot_3d(mesh_gt, mesh_pred_unposed, faces): mesh_pred_unposed[:,2] += 1 fig = go.Figure(data=[ @@ -105,7 +108,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces): return fig -# @tf.function +@tf.function def train_step(batch): with tf.GradientTape() as tape: pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=True) @@ -115,8 +118,12 @@ def train_step(batch): mesh_orig = batch['mesh_unposed'] l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed' + high_resolution_modifier] - pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis]) )), axis=-1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2, power_factors=(0.5, 0.2, 0.2, 0.1))) - # ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2)) + + if low_res == 32: + ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2)) + else: + ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2, power_factors=(0.5, 0.2, 0.2, 0.1))) + loss = 1e3 * (1 - ms_ssim) metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) @@ -125,7 +132,6 @@ def train_step(batch): return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim } - @tf.function def valid_step(batch): pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False) @@ -135,15 +141,18 @@ def valid_step(batch): mesh_orig = batch['mesh_unposed'] l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed' + high_resolution_modifier] - pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis]) )), axis=-1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2, power_factors=(0.5, 0.2, 0.2, 0.1))) - # ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2)) + + if low_res == 32: + ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2)) + else: + ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2, power_factors=(0.5, 0.2, 0.2, 0.1))) + loss = 1e3 * (1 - ms_ssim) metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim } - def plot_everything(batch: dict, type: str): pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False) @@ -168,7 +177,6 @@ def plot_everything(batch: dict, type: str): return { **figs, **{ type + '/images': imgs} } - def train(): valid_iterator = iter(dataset_valid) @@ -204,8 +212,8 @@ def train(): step += 1 if __name__ == '__main__': - # init_wandb() - allow_memory_growth() + init_wandb() + # allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) @@ -214,7 +222,6 @@ if __name__ == '__main__': # uv_mask_body = np.load('custom_utils/mask_256x512.npy') # uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2) - low_res = 32 high_res = 512 uv_mask_body = tf.ones((low_res, low_res//2)) low_resolution_modifier = '_' + str(low_res) @@ -234,8 +241,8 @@ if __name__ == '__main__': 'uv_mask' + low_resolution_modifier: (low_res, low_res//2,), } - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train', 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/train', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + dataset_train = factory.get_dataset(folder='/home/workspace/data/dataset/local-cache/train_skirt', 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/data/dataset/local-cache/train_skirt', 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/valid_skirt', 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/valid_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) @@ -244,14 +251,18 @@ if __name__ == '__main__': # 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 low_res == 32 or low_res == 64: + model = Upscaler(low_res=low_res, kernel_size=3) + elif low_res == 128 or low_res == 256: + model = Upscaler(low_res=low_res, kernel_size=5) + else: + model = Upscaler(low_res=low_res, kernel_size=7) - model = Upscaler(low_res=low_res, kernel_size=3) optimizer = tf.optimizers.Adam(learning_rate=0.0001) folder = 'models/ae_logs/' ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=20) - load_model(model, 'checkpoints/32_512_4') + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + name, max_to_keep=20) train() \ No newline at end of file