From d93678ee1a1a4b020fb0963b8bf14af40fc9c056 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 2 Mar 2023 16:13:18 +0000 Subject: [PATCH] prepare cvae_curv --- aaa_helpers.txt | 2 + configs/iccv/config_cvae_curv.py | 120 +++++++++++++++++++++++++++++++ cvae_curv.py | 51 +++++++------ models/cvae_curv_model.py | 46 +++++++----- 4 files changed, 181 insertions(+), 38 deletions(-) create mode 100644 configs/iccv/config_cvae_curv.py diff --git a/aaa_helpers.txt b/aaa_helpers.txt index 836b02e..139809f 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -37,6 +37,8 @@ nohup python cvae.py --name="lhunor_abl_skin_reg" > nohup/iccv/ablation/abl_skin nohup python cvae.py --name="lhunor_abl_skin_reg_skirt" > nohup/iccv/ablation/abl_skin_reg_skirt.out 2>&1 & nohup python cvae.py --name="lhunor_abl_skin_var" > nohup/iccv/ablation/abl_skin_var.out 2>&1 & +nohup python cvae_curv.py --name="lhunor_cvae_curv" > nohup/iccv/ablation/cvae_curv.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 & diff --git a/configs/iccv/config_cvae_curv.py b/configs/iccv/config_cvae_curv.py new file mode 100644 index 0000000..5ce9ad0 --- /dev/null +++ b/configs/iccv/config_cvae_curv.py @@ -0,0 +1,120 @@ +config = dict() + +# 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 AVG GARMENT USED + +config['use_skinning'] = True # and deshaping +config['remove_global_orientation'] = True +config['use_curvature'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_offsets'] = True + +config['use_norm_loss'] = True + +ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False) +assert ablation_count == 1 or ablation_count == 0 + +config['gpu'] = '5' +config['name'] = 'cvae_curv' +config['notes'] = '' +config['tags'] = ['iccv', 'cvae_curv'] + +config['avg_garment_path'] = '' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv +# if config['use_skinning']: +# assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '') + +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision + +config['batch_size'] = 4 +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False + +config['train_dir'] = '/home/workspace/local-data_127/train' +config['valid_dir'] = '/home/workspace/local-data_127/valid' +config['test_dir'] = '/home/workspace/data/dataset/test_ready' + +config['static_model'] = 'checkpoints/iccv/static_iccv_1' +config['body_model'] = 'checkpoints/iccv/body_iccv_1' +config['curvature_model'] = 'checkpoints/iccv/curv_iccv_1' + + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['pretrained_model'] = '' + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.01 + +config['signature_dynamic'] = { 'id': (), 'frame_nr': (), + 'pose_shape': (72 + 11,), + 'uv_body_posed': (512, 256, 3), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'uv_curvature': (512, 256, 7), + # 'mesh_curvature': (14475) + } +config['signature_static'] = { 'uv_static': (512, 256, 3), + 'mesh_mask': (14475,), + 'uv_mask': (512, 256,), + } + + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False +config['bottleneck_scale'] = 2 + +config['cvae'] = dict() +config['cvae']['is_label'] = True +config['cvae']['latent_dim'] = 64 * config['bottleneck_scale'] +config['cvae']['add_noise'] = False +config['cvae']['noise_amount'] = 0.0 + +config['cvae']['kernel_size'] = 7 +config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3] +config['cvae']['encoder_dims'] = [32, 64, 128, 256, config['cvae']['latent_dim']] +config['cvae']['decoder_depths'] = config['cvae']['encoder_depths'] +config['cvae']['decoder_dims'] = config['cvae']['encoder_dims'] + + +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_curv.py b/cvae_curv.py index bc99e83..d17c125 100644 --- a/cvae_curv.py +++ b/cvae_curv.py @@ -1,6 +1,6 @@ import os import sys -from configs.config_cvae_curv import config +from configs.iccv.config_cvae_curv import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -30,7 +30,6 @@ import plotly.express as ex from plotly.subplots import make_subplots from PIL import Image -from tf_utils import check_tf_version, allow_memory_growth # from losses import collision_penalty as collision_penalty from losses import collision_loss as collision_penalty from dataset import DatasetFactory, StaticDataset @@ -38,15 +37,17 @@ from custom_utils.custom_utils import batch_to_dict from custom_utils.smpl import SMPL from models.cvae_curv_model import CVAE from losses import kl_loss +from custom_utils.custom_utils import Accumulator def init_wandb(): - run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb') + 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_curv.py') artifact.add_file('models/cvae_curv_model.py') + artifact.add_file('models/cvae_model.py') artifact.add_file('dataset.py') artifact.add_file('losses.py') artifact.add_file('models/base.py') @@ -55,7 +56,9 @@ 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.')] 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 = UV2MeshLabel(res=256) @@ -150,7 +153,7 @@ def dice_dense(prediction, ground_truth, weight_map=None): return tf.reduce_mean(dice_score) -@tf.function +# @tf.function def train_step(batch): with tf.GradientTape() as tape: @@ -376,7 +379,9 @@ if __name__ == '__main__': valid_iterator = iter(dataset_valid) ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20) + 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'], iteration=config['pretrained_model_iteration']) @@ -386,29 +391,35 @@ 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) - 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 % 25 == 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) + train_accumulator.update(train_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 % 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() if step % 10000 == 0: manager.save(checkpoint_number=step) - if step == 0: - to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } if step > 40000: g_balancing_coeff = g_balancing_coeff + 1 / 20000 - to_log = { **to_log, **{ 'train/g_balancing_coeff': g_balancing_coeff } } - - wandb.log(to_log, step = step) - print(step) step += 1 \ No newline at end of file diff --git a/models/cvae_curv_model.py b/models/cvae_curv_model.py index c7e8878..7c533a1 100644 --- a/models/cvae_curv_model.py +++ b/models/cvae_curv_model.py @@ -8,7 +8,7 @@ import scipy as sp # import plotly.express as ex from custom_utils.custom_utils import load_model -from models.base import AEEncoder, VariationalEncoder, Decoder +from models.base import AEEncoderNew, VariationalEncoderNew, DecoderNew @@ -23,11 +23,11 @@ class CVAE(tf.keras.Model): self.latent_dim = config['cvae']['latent_dim'] self.size = config['resolution'] - self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size']) + self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation']) if config['static_model'] != '': load_model(self.static_enc, config['static_model'] + '_enc') - self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size']) + self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size']) if config['body_model'] != '': load_model(self.body_enc, config['body_model'] + '_enc') @@ -40,9 +40,10 @@ class CVAE(tf.keras.Model): tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dense(self.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']) - self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'], is_label=config['cvae']['is_label']) + self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'], latent_dim=config['cvae']['latent_dim']) + self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'], is_label=config['cvae']['is_label']) + self.reshape = tf.keras.layers.Reshape((8, 4, config['cvae']['latent_dim'])) self.add_noise = config['cvae']['add_noise'] self.noise_amount = config['cvae']['noise_amount'] @@ -89,12 +90,12 @@ class CVAE(tf.keras.Model): x = self.decoder(z, training=False) - x = x * uv_mask[:,:,:,tf.newaxis] + # x = x * uv_mask[:,:,:,tf.newaxis] - background = tf.ones((uv_template.shape[0],512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis]) - x = tf.concat((background, x), axis=-1) + # background = tf.ones((uv_template.shape[0],512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis]) + # x = tf.concat((background, x), axis=-1) - x = tf.math.softmax(x) + # x = tf.math.softmax(x) return x @@ -112,9 +113,17 @@ class CVAE(tf.keras.Model): return x - def reparameterize(self, mean, logvar): - eps = tf.random.normal(shape=mean.shape) - return eps * tf.exp(logvar * .5) + mean + def reparameterize(self, mean, logvar, training, fix_noise=False): + if training == True: + if not fix_noise: + eps = tf.random.normal(shape=mean.shape) + else: + eps = tf.random.normal(shape=(1, mean.shape[1])) + eps = tf.tile(eps, [mean.shape[0],1]) + out = eps * tf.exp(logvar * .5) + mean + else: + out = tf.exp(logvar * .5) + mean + return out def call(self, inputs, training): @@ -123,18 +132,19 @@ class CVAE(tf.keras.Model): x = tf.concat([uv_template, uv_body_posed, uv_curv], axis=-1) mean, logvar = self.encoder(x) - z = self.reparameterize(mean, logvar) + z = self.reparameterize(mean, logvar, training=training) + z = self.reshape(z) z_cond = self.encode_conds(uv_template, uv_body_posed, training=training) z = tf.concat([z, z_cond], axis=-1) - x = self.decoder(z, training=training) + x = self.decoder([z, uv_mask], training=training) - x = x * uv_mask[:,:,:,tf.newaxis] + # x = x * uv_mask[:,:,:,tf.newaxis] - background = tf.ones((4,512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis]) - x = tf.concat((background, x), axis=-1) + # background = tf.ones((4,512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis]) + # x = tf.concat((background, x), axis=-1) - x = tf.math.softmax(x) + # x = tf.math.softmax(x) return x, mean, logvar