From ec51e5e7cb1f3aee1a45b6b3879553309507bcc3 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 2 Mar 2023 17:40:14 +0000 Subject: [PATCH] resnet pyramid --- configs/iccv/config_abl_pyr_32_res.py | 178 ++++++++++++++++++++++++++ cvae_pyramid.py | 52 ++++---- models/base.py | 4 +- models/pyramid_level.py | 6 +- models/pyramid_model_optimized.py | 6 +- models/resnet.py | 69 ++++++---- models/upscaler.py | 10 +- 7 files changed, 261 insertions(+), 64 deletions(-) create mode 100644 configs/iccv/config_abl_pyr_32_res.py diff --git a/configs/iccv/config_abl_pyr_32_res.py b/configs/iccv/config_abl_pyr_32_res.py new file mode 100644 index 0000000..5e4dc4c --- /dev/null +++ b/configs/iccv/config_abl_pyr_32_res.py @@ -0,0 +1,178 @@ +config = dict() + +config['gpu'] = '5' +config['name'] = 'iccv_pyr_32_res' +config['notes'] = '' +config['tags'] = ['pyr', 'iccv', 'resnet'] + +# 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['use_curvature'] = 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_127/train' +config['valid_dir'] = '/home/workspace/local-data_127/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['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/post/static_resnet_3' +config['pyramid']['body_model'] = 'checkpoints/post/body_resnet_3' +config['pyramid']['curvature_model'] = 'checkpoints/post/curv_resnet_2' +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['bottleneck_size'] = 64 * 2 + level['layer_params'] = [7] + 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']['activation'] = 'tanh' +# config['cond_ae']['layer_params'] = [3, 4, 23, 4, 3] +config['cond_ae_static']['layer_params'] = [3, 7, 15, 7, 3] + +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']['activation'] = 'tanh' +config['cond_ae_curv']['layer_params'] = [3, 7, 15, 7, 3] + +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']['activation'] = 'tanh' +config['cond_ae_body']['layer_params'] = [3, 7, 15, 7, 3] \ No newline at end of file diff --git a/cvae_pyramid.py b/cvae_pyramid.py index cc4cd55..7226078 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_res 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') @@ -208,7 +209,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) @@ -790,7 +791,9 @@ 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/iccv/' + config['name'], max_to_keep=20) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() # load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000') @@ -846,31 +849,32 @@ 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 % 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) + 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 % 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) + train_accumulator.reset() + valid_accumulator.reset() 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 balancing_activated_step == 0: - balancing_activated_step = step - balancing_coeff = balancing_coeff + 1 / 20000 - - - to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } - - wandb.log(to_log, step = step) print(step) diff --git a/models/base.py b/models/base.py index 1b767c6..f8f4cb3 100644 --- a/models/base.py +++ b/models/base.py @@ -113,10 +113,10 @@ class Decoder(tf.keras.Model): class VariationalEncoderNew(tf.keras.Model): - def __init__(self, features): + def __init__(self, layer_params, features): super(VariationalEncoderNew, self).__init__() - self.encoder = resnet_18(features) + self.encoder = resnet_18(layer_params, features) self.mean_branch = tf.keras.Sequential([ tf.keras.layers.Conv2D(filters=features, diff --git a/models/pyramid_level.py b/models/pyramid_level.py index 9cf7322..3e2fd57 100644 --- a/models/pyramid_level.py +++ b/models/pyramid_level.py @@ -102,8 +102,8 @@ 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.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size']) + self.encoder = VariationalEncoderNew(config['layer_params'], config['bottleneck_size']) + self.decoder = DecoderNew(config['layer_params'], config['bottleneck_size']) self.upscale_enhancer = None if config['upscale_enhancer'] != '': @@ -117,7 +117,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['bottleneck_size'])) def enhance(self, low, mask, training=None): diff --git a/models/pyramid_model_optimized.py b/models/pyramid_model_optimized.py index 5301872..5c6ba02 100644 --- a/models/pyramid_model_optimized.py +++ b/models/pyramid_model_optimized.py @@ -256,15 +256,15 @@ class PyramidNew(tf.keras.Model): def __init__(self, config): super(PyramidNew, self).__init__() - 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']) + self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['layer_params'], config['cond_ae_static']['bottleneck_size'], config['cond_ae_static']['activation']) if config['pyramid']['static_model'] != '': load_model(self.static_enc, config['pyramid']['static_model'] + '_enc') - self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size']) + self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['layer_params'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_static']['activation']) if config['pyramid']['body_model'] != '': load_model(self.body_enc, config['pyramid']['body_model'] + '_enc') - self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation']) + self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['layer_params'], config['cond_ae_curv']['bottleneck_size'], config['cond_ae_curv']['activation']) if config['pyramid']['curvature_model'] != '': load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc') diff --git a/models/resnet.py b/models/resnet.py index ef72412..384eb99 100644 --- a/models/resnet.py +++ b/models/resnet.py @@ -131,9 +131,10 @@ class ResNetTypeI(tf.keras.Model): super(ResNetTypeI, self).__init__() self.reversed = reversed self.labels = labels + self.layer_params = layer_params if not reversed: - self.conv1 = tf.keras.layers.Conv2D(filters=32, + self.conv1 = tf.keras.layers.Conv2D(filters=32 if len(layer_params) > 1 else features, kernel_size=(7, 7), strides=2, padding="same") @@ -144,31 +145,36 @@ class ResNetTypeI(tf.keras.Model): else: self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3 if labels is None else labels, kernel_size=(7, 7), - strides=2, + strides=2 if len(layer_params) > 1 else 1, padding="same") self.bn1 = tf.keras.layers.BatchNormalization() self.pool1 = tf.keras.layers.UpSampling2D() - self.layer1 = make_basic_block_layer(filter_num=32, + self.layer1 = make_basic_block_layer(filter_num=32 if len(layer_params) > 1 else (features if not is_cvae else features*2), blocks=layer_params[0], stride=1 if not reversed else 2, upsample=reversed) - self.layer2 = make_basic_block_layer(filter_num=32, - blocks=layer_params[1], - stride=2, - upsample=reversed) - self.layer3 = make_basic_block_layer(filter_num=32, - blocks=layer_params[2], - stride=2, - upsample=reversed) - self.layer4 = make_basic_block_layer(filter_num=64, - blocks=layer_params[3], - stride=2, - upsample=reversed) - self.layer5 = make_basic_block_layer(filter_num=features if not is_cvae else features*2, - blocks=layer_params[4], - stride=2 if not reversed else 1, - upsample=reversed) + if len(layer_params) >= 2: + self.layer2 = make_basic_block_layer(filter_num=32, + blocks=layer_params[1], + stride=2, + upsample=reversed) + if len(layer_params) >= 3: + self.layer3 = make_basic_block_layer(filter_num=32, + blocks=layer_params[2], + stride=2, + upsample=reversed) + if len(layer_params) >= 4: + self.layer4 = make_basic_block_layer(filter_num=64, + blocks=layer_params[3], + stride=2, + upsample=reversed) + + if len(layer_params) >= 5: + self.layer5 = make_basic_block_layer(filter_num=features if not is_cvae else features*2, + blocks=layer_params[4], + stride=2 if not reversed else 1, + upsample=reversed) def call(self, inputs, training=None, mask=None): if not self.reversed: @@ -177,15 +183,24 @@ class ResNetTypeI(tf.keras.Model): x = tf.nn.relu(x) x = self.pool1(x) x = self.layer1(x, training=training) - x = self.layer2(x, training=training) - x = self.layer3(x, training=training) - x = self.layer4(x, training=training) - x = self.layer5(x, training=training) + if len(self.layer_params) >= 2: + x = self.layer2(x, training=training) + if len(self.layer_params) >= 3: + x = self.layer3(x, training=training) + if len(self.layer_params) >= 4: + x = self.layer4(x, training=training) + if len(self.layer_params) >= 5: + x = self.layer5(x, training=training) else: - x = self.layer5(inputs, training=training) - x = self.layer4(x, training=training) - x = self.layer3(x, training=training) - x = self.layer2(x, training=training) + x = inputs + if len(self.layer_params) >= 5: + x = self.layer5(x, training=training) + if len(self.layer_params) >= 4: + x = self.layer4(x, training=training) + if len(self.layer_params) >= 2: + x = self.layer3(x, training=training) + if len(self.layer_params) >= 3: + x = self.layer2(x, training=training) x = self.layer1(x, training=training) x = self.pool1(x) diff --git a/models/upscaler.py b/models/upscaler.py index c03e663..1d65cf3 100644 --- a/models/upscaler.py +++ b/models/upscaler.py @@ -131,16 +131,16 @@ 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'), - ]) + ]) def call(self, inputs, training):