From 8c0436cea8945f1507ba556531a9d5ef0521876d Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 26 Jan 2023 11:10:51 +0000 Subject: [PATCH] restore original pyramid model --- configs/config_abl_pyr_32.py | 2 +- configs/config_abl_pyr_32_skirt.py | 12 +++++++---- cvae_pyramid.py | 34 ++++++++++++++++-------------- models/cvae_model.py | 2 +- 4 files changed, 28 insertions(+), 22 deletions(-) diff --git a/configs/config_abl_pyr_32.py b/configs/config_abl_pyr_32.py index 1181afb..a570b83 100644 --- a/configs/config_abl_pyr_32.py +++ b/configs/config_abl_pyr_32.py @@ -1,7 +1,7 @@ config = dict() config['gpu'] = '0' -config['name'] = 'pyr_abl_32' +config['name'] = 'pyr_abl_32_r1' config['notes'] = '' # DON'T FORGET TO RESET SSIM LOSS USAGE diff --git a/configs/config_abl_pyr_32_skirt.py b/configs/config_abl_pyr_32_skirt.py index 3dd2ee5..7fec07d 100644 --- a/configs/config_abl_pyr_32_skirt.py +++ b/configs/config_abl_pyr_32_skirt.py @@ -1,8 +1,8 @@ config = dict() config['gpu'] = '3' -config['name'] = 'pyr_abl_32_skirt_2' -config['notes'] = 'continued from pyr_abl_32_skirt-180k' +config['name'] = 'pyr_abl_32_skirt_r1' +config['notes'] = '' # DON'T FORGET TO RESET SSIM LOSS USAGE # DON'T FORGET TO CHECK MODEL LOADING @@ -24,7 +24,10 @@ config['pyramid']['resolutions'] = [32] config['pyramid']['pretrained_pyramid'] = '' config['pyramid']['final_upscale_enhancer'] = 'checkpoints/32_512_skirt' -config['train_dir'] = '/home/workspace/local-data/train_skirt' +config['pretrained_model'] = '' +config['pretrained_model_iteration'] = '' + +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 @@ -35,7 +38,7 @@ config['pyramid']['static_model'] = 'checkpoints/static_1_skirt' config['pyramid']['body_model'] = 'checkpoints/body_1_skirt_2_2' config['pyramid']['curvature_model'] = 'checkpoints/curv_1_skirt' -config['dataset_in_memory'] = False +config['dataset_in_memory'] = True config['resolution'] = None config['latent_dim'] = 128 @@ -65,6 +68,7 @@ config['signature_static'] = { 'uv_static': (512, 256, 3), config['pyramid']['is_label'] = False config['pyramid']['latent_dim'] = config['latent_dim'] config['pyramid']['add_noise'] = True +config['pyramid']['noise_amount'] = 0.001 config['pyramid']['loss_weights'] = [1,1,1,1,1] diff --git a/cvae_pyramid.py b/cvae_pyramid.py index cc4cd55..7594a89 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.config_abl_pyr_32_skirt import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -41,7 +41,7 @@ from losses import collision_loss as collision_penalty from custom_utils.custom_utils import batch_to_dict, load_model from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL -from models.pyramid_model_optimized import PyramidNew as Pyramid +from models.pyramid_model_optimized import Pyramid from postprocessing import fix_collision @@ -59,9 +59,11 @@ def init_wandb(): artifact.add_file('models/cond_encoder.py') artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') - config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.post.')] + config_file_names = [key for key in sys.modules.keys() if key.startswith('configs')] 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']) @@ -208,7 +210,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col return fig -@tf.function +# @tf.function def train_step(batch): levels_input, conds = group_inputs_for_levels(batch, config) @@ -261,8 +263,8 @@ def train_step(batch): else: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss) - if config['use_norm_loss'] and step > 2000: - reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] + # if config['use_norm_loss'] and step > 2000: + # reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] kl_total += kl @@ -334,8 +336,8 @@ def valid_step(batch): else: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss) - if config['use_norm_loss'] and step > 2000: - reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] + # if config['use_norm_loss'] and step > 2000: + # reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] kl_total += kl @@ -731,7 +733,7 @@ def speed_step(levels_input, conds): if __name__ == '__main__': - init_wandb() + # init_wandb() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str, default='lhunor') args = vars(parser.parse_args()) @@ -790,7 +792,7 @@ if __name__ == '__main__': # generate_results() ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20) # load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000') @@ -807,18 +809,18 @@ if __name__ == '__main__': model.static_enc.trainable = False model.body_enc.trainable = False model.curv_enc.trainable = False - + model.cond_encoder.trainable = False # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer']) # model.levels[0].upscale_enhancer.trainable = False - for i in range(len(model.levels)): + for i in range(len(model.levels)-1): load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer']) model.levels[i].upscale_enhancer.trainable = False - # if config['pyramid']['resolutions'][-1] != 512: - # load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) - # model.upscale_enhancer.trainable = False + if config['pyramid']['resolutions'][-1] != 512: + load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) + model.upscale_enhancer.trainable = False for i in range(len(model.levels)-1): model.levels[i].trainable = False diff --git a/models/cvae_model.py b/models/cvae_model.py index 3ed4571..f851741 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -43,7 +43,7 @@ class CVAE(tf.keras.Model): self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'), tf.keras.layers.Dense(256, activation='relu'), tf.keras.layers.Dense(128, activation='relu'), - tf.keras.layers.Lambda(lambda x: x / 1000), + # tf.keras.layers.Lambda(lambda x: x / 1000), tf.keras.layers.Dense(config['latent_dim'], activation='tanh')]) self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])