diff --git a/aaa_helpers.txt b/aaa_helpers.txt index b6a7d8e..d6d651a 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -12,6 +12,8 @@ nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_skirt" > nohup/iccv 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 & + skirts!!!! pyramid next week diff --git a/configs/iccv/config_abl_all.py b/configs/iccv/config_abl_all.py index 1ec8467..0ed4371 100644 --- a/configs/iccv/config_abl_all.py +++ b/configs/iccv/config_abl_all.py @@ -18,7 +18,7 @@ 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'] = '0' +config['gpu'] = '3' config['name'] = 'abl_all' config['notes'] = '' @@ -38,9 +38,9 @@ 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['static_model'] = 'checkpoints/static_1' -config['body_model'] = 'checkpoints/body_1' -config['curvature_model'] = 'checkpoints/curv_1' +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 @@ -78,40 +78,47 @@ else: config['upscale_enhancer'] = '' config['freeze_encoders'] = True config['use_one_hot'] = False +config['bottleneck_scale'] = 2 config['cvae'] = dict() config['cvae']['is_label'] = False -config['cvae']['latent_dim'] = config['latent_dim'] -config['cvae']['add_noise'] = True -config['cvae']['noise_amount'] = 0.025 +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, 9, 3] -config['cvae']['encoder_dims'] = [32, 64, 128, 256] -config['cvae']['decoder_input'] = (16, 8, 80) -config['cvae']['decoder_depths'] = [3, 3, 9, 3] -config['cvae']['decoder_dims'] = [256, 128, 64, 16] +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'] = dict() -config['cond_ae']['is_label'] = False -config['cond_ae']['upscale_output'] = False -config['cond_ae']['latent_dim'] = config['latent_dim'] -config['cond_ae']['kernel_size'] = 7 -config['cond_ae']['bottleneck_size'] = 64 -config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3] -config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']] -config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size']) -config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3] -config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16] +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']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['is_label'] = False config['cond_ae_body']['kernel_size'] = 7 -config['cond_ae_body']['upscale_output'] = True -config['cond_ae_body']['bottleneck_size'] = 512 -config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3] -config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']] -config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size']) -config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3] -config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 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'] +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 f79cdd1..8471adb 100644 --- a/cvae.py +++ b/cvae.py @@ -1,6 +1,6 @@ import os import sys -from configs.config_abl_all_post import config +from configs.iccv.config_abl_all import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -45,7 +45,7 @@ from models.cvae_model import CVAENew as CVAE def init_wandb(): - run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/tmp', 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.py') artifact.add_file('models/cvae_model.py') @@ -75,7 +75,7 @@ def kl_loss(mean, logvar): # 3d loss on posed !!!!!!!!!!!! -# @tf.function +@tf.function def train_step(batch): with tf.GradientTape() as tape: @@ -478,7 +478,7 @@ def test(): if __name__ == '__main__': - # init_wandb() + init_wandb() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) @@ -528,7 +528,7 @@ if __name__ == '__main__': # test() 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/iccv/' + config['name'], max_to_keep=20) if config['pretrained_model'] != '': @@ -543,9 +543,10 @@ if __name__ == '__main__': train_loss = train_step(batch) - 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 % 10 == 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) to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } if step % 5000 == 0 or (step < 5000 and step % 500 == 0): diff --git a/models/base.py b/models/base.py index e9cfd67..2058865 100644 --- a/models/base.py +++ b/models/base.py @@ -111,18 +111,18 @@ class Decoder(tf.keras.Model): class VariationalEncoderNew(tf.keras.Model): - def __init__(self, depths, dims, kernel_size=7): + def __init__(self, depths, dims, kernel_size=7, latent_dim=64): super(VariationalEncoderNew, self).__init__() self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size) self.mean_branch = tf.keras.Sequential( - [Block(dim=64, drop_path=0, kernel_size=kernel_size, + [Block(dim=latent_dim, drop_path=0, kernel_size=kernel_size, layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] ) self.var_branch = tf.keras.Sequential( - [Block(dim=64, drop_path=0, kernel_size=kernel_size, + [Block(dim=latent_dim, drop_path=0, kernel_size=kernel_size, layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] ) diff --git a/models/cvae_model.py b/models/cvae_model.py index 3ed4571..3e25769 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -179,10 +179,10 @@ class CVAENew(tf.keras.Model): if config['use_curvature']: self.curv_enc.trainable = False - self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size']) + 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']) - self.reshape = tf.keras.layers.Reshape((8, 4, 64)) + self.reshape = tf.keras.layers.Reshape((8, 4, config['cvae']['latent_dim'])) if config['avg_garment_path'] != '': self.avg_garment = np.load(config['avg_garment_path'])