new weight trials, 10_3 didn't run

This reverts commit ac27319860.

aaa_helpers.txt

@@ -1,4 +1,4 @@
-nohup python cvae.py --name="lhunor_weights_test_6" > nohup/iccv/weights_test_6.out 2>&1 &
+nohup python cvae.py --name="lhunor_weights_test_10_2" > nohup/iccv/weights_test_10_2.out 2>&1 &
 nohup python evaluate_ae.py --name iccv_weights_test_6 > nohup/iccv/evaluate_iccv_weights_test_6.out 2>&1 &
 nohup python calculate_error.py --name iccv_weights_test_6 > nohup/iccv/error_iccv_weights_test_6.out 2>&1 &
 

configs/config_abl_all_post.py

@@ -19,13 +19,13 @@ 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'] = '5'
+config['gpu'] = '6'
-config['name'] = 'iccv_weights_test_8'
+config['name'] = 'iccv_weights_test_10_3'
-config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, changed convnext reverse'
+config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, add gnn but smaller, same as 10 but training from zero and no avg garment'
 
-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
+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']:
+# if config['use_skinning']:
-  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+#   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
 
@@ -91,9 +91,9 @@ config['cvae']['noise_amount'] = 0.00001
 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'] = [3, 9, 3, 3, 3, 3]
+config['cvae']['decoder_depths'] = config['cvae']['encoder_depths']
-config['cvae']['decoder_dims'] = [3, 32, 128, 128, 256, 512]
+config['cvae']['decoder_dims'] = config['cvae']['encoder_dims']
-config['cvae']['decoder_weights_depths'] = [3, 9, 3, 3, 3]
+config['cvae']['decoder_weights_depths'] = [3, 3, 3, 5, 3]
 config['cvae']['decoder_weights_dims'] = [24, 32, 64, 128, 256]
 
 config['cond_ae_static'] = dict()

convnext/models/convnext_tf.py

@@ -225,28 +225,26 @@ class ConvNeXtReverse(tf.keras.Model):
                  drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., learn_weights=False,
                  ):
         super().__init__()
-        self.upsample_layers = [] 
+        self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
-
         stem = tf.keras.Sequential([
-            LayerNormalization(epsilon=1e-6),
+            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
-            Conv2DTranspose(dims[0], kernel_size=2, strides=2, padding='same'),
+            # UpSampling2D((4, 4), interpolation='nearest'),
-        ])
+            # Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
+            # Dense(dims[0]),
+            LayerNormalization(epsilon=1e-6)]
+        )
         self.upsample_layers.append(stem)
-
+        for i in range(len(depths) - 1):
-        for i in range(len(depths) - 2):
             downsample_layer = tf.keras.Sequential([
                 LayerNormalization(epsilon=1e-6),
-                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same'),
+                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
+                # UpSampling2D((2,2), interpolation='nearest'),
+                # Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
+                # Dense(dims[i+1]),
             ])
             self.upsample_layers.append(downsample_layer)
 
-        downsample_layer = tf.keras.Sequential([
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
-            Conv2DTranspose(dims[i+2], kernel_size=2, strides=2, padding='same'),
-            LayerNormalization(epsilon=1e-6),
-        ])
-        self.upsample_layers.append(downsample_layer)
-
-        self.stages = []  
         dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
         cur = 0
         for i in range(len(depths)):
@@ -259,9 +257,9 @@ class ConvNeXtReverse(tf.keras.Model):
 
         if not learn_weights:
             self.last = Dense(3, activation='tanh')
-            # self.last = Activation('tanh')
         else:
             self.last = Activation('softmax')
+            # self.last = lambda x, training: x
             # self.last = Dense(24, activation='softmax')
 
 

cvae.py

@@ -42,6 +42,7 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mes
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
 from models.cvae_model import CVAENew as CVAE
+from models.gnn import Smoother
 
 
 def init_wandb():
@@ -92,7 +93,9 @@ def train_step(batch):
         pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
 
         mesh_pred = uv2mesh(pred)
+        mesh_pred = smoother(mesh_pred, training=True)
         mesh_weigths = uv2mesh_label(weights)
+        mesh_weigths = smoother_weights(mesh_weigths, training=True)
 
         if config['use_skinning']:
             shape = batch['pose_shape'][:, 72:]
@@ -137,8 +140,8 @@ def train_step(batch):
 
     # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_posed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    gradients = tape.gradient(loss, model.trainable_variables)
+    gradients = tape.gradient(loss, model.trainable_variables + smoother.trainable_variables + smoother_weights.trainable_variables)
-    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables + smoother.trainable_variables + smoother_weights.trainable_variables))
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
              'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
              'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
@@ -157,7 +160,9 @@ def valid_step(batch):
     pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
 
     mesh_pred = uv2mesh(pred)
+    mesh_pred = smoother(mesh_pred, training=False)
     mesh_weigths = uv2mesh_label(weights)
+    mesh_weigths = smoother_weights(mesh_weigths, training=False)
 
     if config['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
@@ -273,8 +278,10 @@ def plot_everything(batch: dict, type: str):
 
     pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
     
-    mesh_pred = uv2mesh(pred).numpy()
+    mesh_pred = uv2mesh(pred)
+    mesh_pred = smoother(mesh_pred, training=False).numpy()
     mesh_weigths = uv2mesh_label(weights)
+    mesh_weigths = smoother_weights(mesh_weigths, training=False)
     
     orig_img = uv_orig
     uv_static_img = batch['uv_static']
@@ -282,7 +289,10 @@ def plot_everything(batch: dict, type: str):
     pred_img = pred
     img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
-    
+    # min max normalization of weights
+    weights = (weights - tf.reduce_min(weights, axis=[1,2,3], keepdims=True)) / (tf.reduce_max(weights, axis=[1,2,3], keepdims=True) - tf.reduce_min(weights, axis=[1,2,3], keepdims=True))
+    weights_imgs = [wandb.Image(np.concatenate((np.concatenate([weights[b,:,:,i].numpy() for i in range(12)], axis=-1), np.concatenate([weights[b,:,:,i].numpy() for i in range(12,24)], axis=-1)), axis=0)) for b in range(4)]
+
     F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
 
@@ -332,7 +342,7 @@ def plot_everything(batch: dict, type: str):
     mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0))
     logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0))
 
-    return { **figs, **figs_coll, **figs_weights, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+    return { **figs, **figs_coll, **figs_weights, **{ type + '/images': imgs, type + '/weight_images': weights_imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
 def sample(folder): 
@@ -547,6 +557,8 @@ if __name__ == '__main__':
                                         is_cape=config['is_cape'])
 
     model = CVAE(config)
+    smoother = Smoother(features=24, activation_middle='tanh', activation_final='tanh', learn_weights=False)
+    smoother_weights = Smoother(features=3, activation_middle=None, activation_final='softmax', learn_weights=True)
 
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
     balancing_activated_step = 0
@@ -571,9 +583,10 @@ if __name__ == '__main__':
 
         train_loss = train_step(batch)
 
-        data_valid = valid_iterator.get_next()
+        if step % 25 == 0:
-        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
+            data_valid = valid_iterator.get_next()
-        valid_loss = valid_step(data_valid)
+            data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
+            valid_loss = valid_step(data_valid)
 
         to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
         if step % 5000 == 0 or (step < 5000 and step % 500 == 0):

models/gnn.py

@@ -0,0 +1,34 @@
+import os
+import sys
+
+import scipy.io as sio
+import tensorflow as tf
+from spektral.layers import GCNConv, GeneralConv
+from spektral.utils import sp_matrix_to_sp_tensor
+from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
+                                     GlobalAveragePooling2D, Layer,
+                                     LayerNormalization, Flatten, Reshape, UpSampling2D)
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+class Smoother(tf.keras.Model):
+    def __init__(self, features=64, activation_middle='tanh', activation_final='tanh', learn_weights=False):
+        super().__init__()
+        self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo())
+        self.conv1 = GCNConv(features, activation=activation_middle, use_bias=False)
+        self.conv2 = GCNConv(features, activation=activation_middle, use_bias=False)
+        # self.conv3 = GCNConv(features, activation=activation_middle, use_bias=False)
+        # self.conv4 = GCNConv(64, activation='tanh', use_bias=False)
+        # self.conv5 = GCNConv(64, activation='tanh', use_bias=False)
+        self.conv6 = GCNConv(3 if not learn_weights else 24, activation=activation_final, use_bias=False)
+
+    def call(self, inputs, training=None):
+        x = inputs
+        x = self.conv1([x, self.a], training=training)
+        x = self.conv2([x, self.a], training=training)
+        # x = self.conv3([x, self.a], training=training)
+        # x = self.conv4([x, self.a], training=training)
+        # x = self.conv5([x, self.a], training=training)
+        x = self.conv6([x, self.a], training=training)
+        return x