creating resnet101

This commit is contained in:
2023-02-21 22:17:40 +00:00
parent e302deab7e
commit 8447b7c505
7 changed files with 370 additions and 57 deletions
+3 -4
View File
@@ -1,6 +1,6 @@
import os import os
import sys 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'] os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf 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): def train_step(batch):
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
pred = model([batch['uv_static'], batch['uv_mask']], training=True) pred = model([batch['uv_static'], batch['uv_mask']], training=True)
@@ -189,8 +189,7 @@ def test():
if __name__ == '__main__': if __name__ == '__main__':
init_wandb() # init_wandb()
# allow_memory_growth()
parser = argparse.ArgumentParser(description='') parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str) parser.add_argument('--name', type=str)
args = vars(parser.parse_args()) args = vars(parser.parse_args())
+126
View File
@@ -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'
+51
View File
@@ -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'],),
}
+50
View File
@@ -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,),
}
+51
View File
@@ -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,),
}
+6 -8
View File
@@ -1,6 +1,6 @@
import tensorflow as tf import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels 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): class Bottleneck(tf.keras.Model):
@@ -116,7 +116,7 @@ class VariationalEncoderNew(tf.keras.Model):
def __init__(self, features): def __init__(self, features):
super(VariationalEncoderNew, self).__init__() super(VariationalEncoderNew, self).__init__()
self.encoder = resnet_18(features) self.encoder = resnet_101(features)
self.mean_branch = tf.keras.Sequential([ self.mean_branch = tf.keras.Sequential([
tf.keras.layers.Conv2D(filters=features, tf.keras.layers.Conv2D(filters=features,
@@ -166,7 +166,7 @@ class AEEncoderNew(tf.keras.Model):
def __init__(self, features, activation='tanh'): def __init__(self, features, activation='tanh'):
super(AEEncoderNew, self).__init__() super(AEEncoderNew, self).__init__()
self.encoder = resnet_18(features) self.encoder = resnet_101(features)
self.act = tf.keras.layers.Activation(activation) self.act = tf.keras.layers.Activation(activation)
@@ -187,17 +187,15 @@ class AEEncoderNew(tf.keras.Model):
class DecoderNew(tf.keras.Model): class DecoderNew(tf.keras.Model):
def __init__(self, features, is_label=False, is_cvae=False): def __init__(self, features, is_label=False, is_cvae=False, learn_weights=False):
def __init__(self, depths, dims, kernel_size=7, is_label=False, learn_weights=False):
super(DecoderNew, self).__init__() super(DecoderNew, self).__init__()
self.is_label = is_label self.is_label = is_label
if not is_label: if not is_label:
self.decoder = resnet_18(features, reversed=True, is_cvae=is_cvae) self.decoder = resnet_101(features, reversed=True, is_cvae=is_cvae)
self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size, learn_weights=learn_weights)
else: 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): def call(self, inputs, training):
+83 -45
View File
@@ -54,17 +54,23 @@ class BasicBlock(tf.keras.layers.Layer):
class BottleNeck(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__() super(BottleNeck, self).__init__()
self.conv1 = tf.keras.layers.Conv2D(filters=filter_num, self.conv1 = tf.keras.layers.Conv2D(filters=filter_num,
kernel_size=(1, 1), kernel_size=(1, 1),
strides=1, strides=1,
padding='same') padding='same')
self.bn1 = tf.keras.layers.BatchNormalization() self.bn1 = tf.keras.layers.BatchNormalization()
self.conv2 = tf.keras.layers.Conv2D(filters=filter_num, if not upsample:
kernel_size=(3, 3), self.conv2 = tf.keras.layers.Conv2D(filters=filter_num,
strides=stride, kernel_size=(3, 3),
padding='same') 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.bn2 = tf.keras.layers.BatchNormalization()
self.conv3 = tf.keras.layers.Conv2D(filters=filter_num * 4, self.conv3 = tf.keras.layers.Conv2D(filters=filter_num * 4,
kernel_size=(1, 1), kernel_size=(1, 1),
@@ -72,14 +78,19 @@ class BottleNeck(tf.keras.layers.Layer):
padding='same') padding='same')
self.bn3 = tf.keras.layers.BatchNormalization() self.bn3 = tf.keras.layers.BatchNormalization()
self.downsample = tf.keras.Sequential() self.updownsample = tf.keras.Sequential()
self.downsample.add(tf.keras.layers.Conv2D(filters=filter_num * 4, if not upsample:
kernel_size=(1, 1), self.updownsample.add(tf.keras.layers.Conv2D(filters=filter_num * 4,
strides=stride)) kernel_size=(1, 1),
self.downsample.add(tf.keras.layers.BatchNormalization()) 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): def call(self, inputs, training=None, **kwargs):
residual = self.downsample(inputs) residual = self.updownsample(inputs)
x = self.conv1(inputs) x = self.conv1(inputs)
x = self.bn1(x, training=training) 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 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 = 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): for _ in range(1, blocks):
res_block.add(BottleNeck(filter_num, stride=1)) res_block.add(BottleNeck(filter_num, stride=1))
@@ -136,9 +147,6 @@ class ResNetTypeI(tf.keras.Model):
strides=2, strides=2,
padding="same") padding="same")
self.bn1 = tf.keras.layers.BatchNormalization() 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.pool1 = tf.keras.layers.UpSampling2D()
self.layer1 = make_basic_block_layer(filter_num=64, self.layer1 = make_basic_block_layer(filter_num=64,
@@ -193,45 +201,75 @@ class ResNetTypeI(tf.keras.Model):
class ResNetTypeII(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__() super(ResNetTypeII, self).__init__()
self.conv1 = tf.keras.layers.Conv2D(filters=64, self.reversed = reversed
kernel_size=(7, 7), self.labels = labels
strides=2,
padding="same") if not reversed:
self.bn1 = tf.keras.layers.BatchNormalization() self.conv1 = tf.keras.layers.Conv2D(filters=64,
self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3), kernel_size=(7, 7),
strides=2, strides=2,
padding="same") 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, 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, self.layer2 = make_bottleneck_layer(filter_num=128,
blocks=layer_params[1], blocks=layer_params[1],
stride=2) stride=2,
upsample=reversed)
self.layer3 = make_bottleneck_layer(filter_num=256, self.layer3 = make_bottleneck_layer(filter_num=256,
blocks=layer_params[2], blocks=layer_params[2],
stride=2) stride=2,
upsample=reversed)
self.layer4 = make_bottleneck_layer(filter_num=512, self.layer4 = make_bottleneck_layer(filter_num=512,
blocks=layer_params[3], blocks=layer_params[3],
stride=2) stride=2,
upsample=reversed)
self.avgpool = tf.keras.layers.GlobalAveragePooling2D() self.layer5 = make_bottleneck_layer(filter_num=features if not is_cvae else features*2,
self.fc = tf.keras.layers.Dense(units=NUM_CLASSES, activation=tf.keras.activations.softmax) blocks=layer_params[4],
stride=2 if not reversed else 1,
upsample=reversed)
def call(self, inputs, training=None, mask=None): def call(self, inputs, training=None, mask=None):
x = self.conv1(inputs) if not self.reversed:
x = self.bn1(x, training=training) x = self.conv1(inputs)
x = tf.nn.relu(x) x = self.bn1(x, training=training)
x = self.pool1(x) x = tf.nn.relu(x)
x = self.layer1(x, training=training) x = self.pool1(x)
x = self.layer2(x, training=training) x = self.layer1(x, training=training)
x = self.layer3(x, training=training) x = self.layer2(x, training=training)
x = self.layer4(x, training=training) x = self.layer3(x, training=training)
x = self.avgpool(x) x = self.layer4(x, training=training)
output = self.fc(x) 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): 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]) return ResNetTypeII(layer_params=[3, 4, 6, 3])
def resnet_101(): def resnet_101(features, reversed=False, labels=None, is_cvae=False):
return ResNetTypeII(layer_params=[3, 4, 23, 3]) return ResNetTypeII(layer_params=[3, 4, 23, 3, 3], features=features, reversed=reversed, labels=labels, is_cvae=is_cvae)
def resnet_152(): def resnet_152():