resnet pyramid

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]

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)
 
+        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)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+
+        if step % 5000 == 0:
+            manager.save(checkpoint_number=step)
 
         print(step)
 

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,

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.encoder = VariationalEncoderNew(config['layer_params'], config['bottleneck_size'])
-        self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_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

@@ -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

@@ -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

@@ -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'),
                     ])