From 8447b7c50553cf9145eb254270d49eaba08505b4 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 21 Feb 2023 22:17:40 +0000 Subject: [PATCH] creating resnet101 --- auto_encoder_static.py | 7 +- configs/iccv/config_abl_all.py | 126 ++++++++++++++++++++++++++++++ configs/iccv/config_ae_body.py | 51 ++++++++++++ configs/iccv/config_ae_curv.py | 50 ++++++++++++ configs/iccv/config_ae_static.py | 51 ++++++++++++ models/base.py | 14 ++-- models/resnet.py | 128 ++++++++++++++++++++----------- 7 files changed, 370 insertions(+), 57 deletions(-) create mode 100644 configs/iccv/config_abl_all.py create mode 100644 configs/iccv/config_ae_body.py create mode 100644 configs/iccv/config_ae_curv.py create mode 100644 configs/iccv/config_ae_static.py diff --git a/auto_encoder_static.py b/auto_encoder_static.py index 87325c3..e0f9a68 100644 --- a/auto_encoder_static.py +++ b/auto_encoder_static.py @@ -1,6 +1,6 @@ import os import sys -from configs.review.config_ae_static import config +from configs.iccv.config_ae_static import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -55,7 +55,7 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_ -@tf.function +# @tf.function def train_step(batch): with tf.GradientTape() as tape: pred = model([batch['uv_static'], batch['uv_mask']], training=True) @@ -189,8 +189,7 @@ def test(): if __name__ == '__main__': - init_wandb() - # allow_memory_growth() + # init_wandb() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) diff --git a/configs/iccv/config_abl_all.py b/configs/iccv/config_abl_all.py new file mode 100644 index 0000000..68985b5 --- /dev/null +++ b/configs/iccv/config_abl_all.py @@ -0,0 +1,126 @@ +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 +config['only_curvatures'] = False + +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'] = '4' +config['name'] = 'abl_all_resnet_r7' +config['notes'] = '' + +config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #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/train' +config['valid_dir'] = '/home/workspace/local-data/valid' +config['test_dir'] = '/home/workspace/data/dataset/test_ready' + +# config['frames_file'] = 'frames_dress_30k.txt' +# config['frames_file_valid'] = 'frames_dress_30k.txt' +# config['frames_file_test'] = 'frames_dress_30k.txt' + +config['static_model'] = 'checkpoints/post/static_r7_3' +config['body_model'] = 'checkpoints/post/body_r7_3' +config['curvature_model'] = 'checkpoints/post/curv_r7_3' + + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + + +config['pretrained_model'] = '' + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.001 + +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,), + 'uv_curvature': (512, 256, 7), + } +config['signature_static'] = { 'uv_static': (512, 256, 3), + 'mesh_mask': (config['smpl_res'],), + 'uv_mask': (512, 256,), + } + +if config['use_skinning']: + config['signature_dynamic']['uv_unposed'] = (512, 256, 3) + config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,) +else: + config['signature_dynamic']['uv_posed'] = (512, 256, 3) + + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False + +config['cvae'] = dict() +config['cvae']['is_label'] = False +config['cvae']['add_noise'] = True +config['cvae']['noise_amount'] = 0.001 + +config['cvae']['kernel_size'] = 7 +config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3] +config['cvae']['encoder_dims'] = [32, 64, 128, 256, 64] +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['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_body.py b/configs/iccv/config_ae_body.py new file mode 100644 index 0000000..5018050 --- /dev/null +++ b/configs/iccv/config_ae_body.py @@ -0,0 +1,51 @@ +config = dict() + + +config['batch_size'] = 4 +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['use_filled'] = False +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +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'] = 16 +config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3] +config['cond_ae']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae']['bottleneck_size']] +config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths'] +config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims'] +config['cond_ae']['activation'] = 'tanh' + + +config['gpu'] = '4' +config['lr'] = 0.0001 +config['name'] = 'body_r7_3' +config['notes'] = 'resnet18' +config['tags'] = ['review_7'] + +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + + +config['signature_dynamic'] = { 'id': (), 'frame_nr': (), + 'pose_shape': (72 + 11,), + 'uv_body_posed': (3, 512, 256,), + 'mesh_body_posed': (config['smpl_res'], 3,), + } +config['signature_static'] = { + 'mesh_mask': (config['smpl_res'],), + } \ No newline at end of file diff --git a/configs/iccv/config_ae_curv.py b/configs/iccv/config_ae_curv.py new file mode 100644 index 0000000..3a54856 --- /dev/null +++ b/configs/iccv/config_ae_curv.py @@ -0,0 +1,50 @@ +config = dict() + + +config['batch_size'] = 4 +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['use_filled'] = False +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['remove_global_orientation'] = True + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['cond_ae'] = dict() +config['cond_ae']['is_label'] = True +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'] = 40 +config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3] +config['cond_ae']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae']['bottleneck_size']] +config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths'] +config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims'] +config['cond_ae']['activation'] = 'tanh' + + +config['gpu'] = '3' +config['lr'] = 0.001 +config['name'] = 'curv_r7_3' +config['notes'] = 'resnet18' +config['tags'] = ['review_7'] + +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + + +config['signature_dynamic'] = { 'id': (), 'frame_nr': (), + 'uv_curvature': (512, 256, 7), + 'mesh_curvature': (config['smpl_res']) + } +config['signature_static'] = { 'mesh_mask': (config['smpl_res'],), + 'uv_mask': (512, 256,), + } \ No newline at end of file diff --git a/configs/iccv/config_ae_static.py b/configs/iccv/config_ae_static.py new file mode 100644 index 0000000..641ff1a --- /dev/null +++ b/configs/iccv/config_ae_static.py @@ -0,0 +1,51 @@ +config = dict() + + +config['batch_size'] = 4 +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['use_filled'] = False +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = False +config['remove_global_orientation'] = True + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +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'] = 8 * 4 +config['cond_ae']['encoder_depths'] = [3, 3, 3, 27, 3] +config['cond_ae']['encoder_dims'] = [8, 64, 128, 128, config['cond_ae']['bottleneck_size']] +config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths'] +config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims'] +config['cond_ae']['activation'] = 'linear' + + +config['gpu'] = '5' +config['lr'] = 0.0005 +config['name'] = 'static_resnet_1' +config['notes'] = 'resnet101 based, one more block' +config['tags'] = ['iccv', 'resnet'] + + + +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + + +config['signature_dynamic'] = { 'id': (), 'frame_nr': (), } +config['signature_static'] = { 'uv_static': (512, 256, 3), + 'mesh_static': (config['smpl_res'], 3,), + 'mesh_mask': (config['smpl_res'],), + 'uv_mask': (512, 256,), + } \ No newline at end of file diff --git a/models/base.py b/models/base.py index 6f0f39b..e9b9c83 100644 --- a/models/base.py +++ b/models/base.py @@ -1,6 +1,6 @@ import tensorflow as tf from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels -from models.resnet import resnet_18, make_basic_block_layer +from models.resnet import resnet_18, make_basic_block_layer, resnet_34, resnet_50, resnet_101, resnet_152 class Bottleneck(tf.keras.Model): @@ -116,7 +116,7 @@ class VariationalEncoderNew(tf.keras.Model): def __init__(self, features): super(VariationalEncoderNew, self).__init__() - self.encoder = resnet_18(features) + self.encoder = resnet_101(features) self.mean_branch = tf.keras.Sequential([ tf.keras.layers.Conv2D(filters=features, @@ -166,7 +166,7 @@ class AEEncoderNew(tf.keras.Model): def __init__(self, features, activation='tanh'): super(AEEncoderNew, self).__init__() - self.encoder = resnet_18(features) + self.encoder = resnet_101(features) self.act = tf.keras.layers.Activation(activation) @@ -187,17 +187,15 @@ class AEEncoderNew(tf.keras.Model): class DecoderNew(tf.keras.Model): - def __init__(self, features, is_label=False, is_cvae=False): - def __init__(self, depths, dims, kernel_size=7, is_label=False, learn_weights=False): + def __init__(self, features, is_label=False, is_cvae=False, learn_weights=False): super(DecoderNew, self).__init__() self.is_label = is_label if not is_label: - self.decoder = resnet_18(features, reversed=True, is_cvae=is_cvae) - self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size, learn_weights=learn_weights) + self.decoder = resnet_101(features, reversed=True, is_cvae=is_cvae) else: - self.decoder = resnet_18(features, reversed=True, labels=6, is_cvae=is_cvae) + self.decoder = resnet_101(features, reversed=True, labels=6, is_cvae=is_cvae) def call(self, inputs, training): diff --git a/models/resnet.py b/models/resnet.py index 4673404..dc47cb9 100644 --- a/models/resnet.py +++ b/models/resnet.py @@ -54,17 +54,23 @@ class BasicBlock(tf.keras.layers.Layer): class BottleNeck(tf.keras.layers.Layer): - def __init__(self, filter_num, stride=1): + def __init__(self, filter_num, stride=1, upsample=False): super(BottleNeck, self).__init__() self.conv1 = tf.keras.layers.Conv2D(filters=filter_num, kernel_size=(1, 1), strides=1, padding='same') self.bn1 = tf.keras.layers.BatchNormalization() - self.conv2 = tf.keras.layers.Conv2D(filters=filter_num, - kernel_size=(3, 3), - strides=stride, - padding='same') + if not upsample: + self.conv2 = tf.keras.layers.Conv2D(filters=filter_num, + kernel_size=(3, 3), + strides=stride, + padding='same') + else: + self.conv2 = tf.keras.layers.Conv2DTranspose(filters=filter_num, + kernel_size=(3, 3), + strides=stride, + padding='same') self.bn2 = tf.keras.layers.BatchNormalization() self.conv3 = tf.keras.layers.Conv2D(filters=filter_num * 4, kernel_size=(1, 1), @@ -72,14 +78,19 @@ class BottleNeck(tf.keras.layers.Layer): padding='same') self.bn3 = tf.keras.layers.BatchNormalization() - self.downsample = tf.keras.Sequential() - self.downsample.add(tf.keras.layers.Conv2D(filters=filter_num * 4, - kernel_size=(1, 1), - strides=stride)) - self.downsample.add(tf.keras.layers.BatchNormalization()) + self.updownsample = tf.keras.Sequential() + if not upsample: + self.updownsample.add(tf.keras.layers.Conv2D(filters=filter_num * 4, + kernel_size=(1, 1), + strides=stride)) + else: + self.updownsample.add(tf.keras.layers.Conv2DTranspose(filters=filter_num * 4, + kernel_size=(1, 1), + strides=stride)) + self.updownsample.add(tf.keras.layers.BatchNormalization()) def call(self, inputs, training=None, **kwargs): - residual = self.downsample(inputs) + residual = self.updownsample(inputs) x = self.conv1(inputs) x = self.bn1(x, training=training) @@ -105,9 +116,9 @@ def make_basic_block_layer(filter_num, blocks, stride=1, upsample=False): return res_block -def make_bottleneck_layer(filter_num, blocks, stride=1): +def make_bottleneck_layer(filter_num, blocks, stride=1, upsample=False): res_block = tf.keras.Sequential() - res_block.add(BottleNeck(filter_num, stride=stride)) + res_block.add(BottleNeck(filter_num, stride=stride, upsample=upsample)) for _ in range(1, blocks): res_block.add(BottleNeck(filter_num, stride=1)) @@ -136,9 +147,6 @@ class ResNetTypeI(tf.keras.Model): strides=2, padding="same") self.bn1 = tf.keras.layers.BatchNormalization() - # self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), - # strides=2, - # padding="same") self.pool1 = tf.keras.layers.UpSampling2D() self.layer1 = make_basic_block_layer(filter_num=64, @@ -193,45 +201,75 @@ class ResNetTypeI(tf.keras.Model): class ResNetTypeII(tf.keras.Model): - def __init__(self, layer_params): + def __init__(self, layer_params, features, reversed=False, labels=None, is_cvae=False): super(ResNetTypeII, self).__init__() - self.conv1 = tf.keras.layers.Conv2D(filters=64, - kernel_size=(7, 7), - strides=2, - padding="same") - self.bn1 = tf.keras.layers.BatchNormalization() - self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), - strides=2, - padding="same") + self.reversed = reversed + self.labels = labels + + if not reversed: + self.conv1 = tf.keras.layers.Conv2D(filters=64, + kernel_size=(7, 7), + strides=2, + padding="same") + self.bn1 = tf.keras.layers.BatchNormalization() + self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), + strides=2, + padding="same") + else: + self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3 if labels is None else labels, + kernel_size=(7, 7), + strides=2, + padding="same") + self.bn1 = tf.keras.layers.BatchNormalization() + self.pool1 = tf.keras.layers.UpSampling2D() self.layer1 = make_bottleneck_layer(filter_num=64, - blocks=layer_params[0]) + blocks=layer_params[0], + stride=1 if not reversed else 2, + upsample=reversed) self.layer2 = make_bottleneck_layer(filter_num=128, blocks=layer_params[1], - stride=2) + stride=2, + upsample=reversed) self.layer3 = make_bottleneck_layer(filter_num=256, blocks=layer_params[2], - stride=2) + stride=2, + upsample=reversed) self.layer4 = make_bottleneck_layer(filter_num=512, blocks=layer_params[3], - stride=2) - - self.avgpool = tf.keras.layers.GlobalAveragePooling2D() - self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, activation=tf.keras.activations.softmax) + stride=2, + upsample=reversed) + self.layer5 = make_bottleneck_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): - x = self.conv1(inputs) - x = self.bn1(x, training=training) - 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.avgpool(x) - output = self.fc(x) + if not self.reversed: + x = self.conv1(inputs) + x = self.bn1(x, training=training) + 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) + 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 = self.layer1(x, training=training) - return output + x = self.pool1(x) + x = self.bn1(x, training=training) + x = self.conv1(x) + if self.labels is None: + x = tf.nn.tanh(x) + + return x def resnet_18(features, reversed=False, labels=None, is_cvae=False): @@ -246,8 +284,8 @@ def resnet_50(): return ResNetTypeII(layer_params=[3, 4, 6, 3]) -def resnet_101(): - return ResNetTypeII(layer_params=[3, 4, 23, 3]) +def resnet_101(features, reversed=False, labels=None, is_cvae=False): + return ResNetTypeII(layer_params=[3, 4, 23, 3, 3], features=features, reversed=reversed, labels=labels, is_cvae=is_cvae) def resnet_152():