luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

resnet pyramid

Browse Source
This commit is contained in:
luni
2023-03-02 17:40:14 +00:00
parent 532f5e12f4
commit ec51e5e7cb
7 changed files with 261 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/iccv/config_abl_pyr_32_res.py
+178
View File
@@ -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]
cvae_pyramid.py
+20 -16
View File
@@ -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)
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):
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 % 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)
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
if step % 5000 == 0:
manager.save(checkpoint_number=step)
print(step)
models/base.py
+2 -2
View File
@@ -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,
models/pyramid_level.py
+3 -3
View File
@@ -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):
models/pyramid_model_optimized.py
+3 -3
View File
@@ -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')
models/resnet.py
+19 -4
View File
@@ -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,27 +145,32 @@ 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)
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,
@@ -177,14 +183,23 @@ class ResNetTypeI(tf.keras.Model):
x = tf.nn.relu(x)
x = self.pool1(x)
x = self.layer1(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 = 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)
models/upscaler.py
+4 -4
View File
@@ -131,13 +131,13 @@ 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'),
])