diff --git a/aaa_helpers.txt b/aaa_helpers.txt index 3f66e3e..ea6153a 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -14,6 +14,7 @@ nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_m" > nohup/icc nohup python cvae.py --name="lhunor_cvae_iccv_gnn_1" > nohup/iccv/cvae_iccv_gnn_1.out 2>&1 & nohup python cvae.py --name="lhunor_abl_all" > nohup/iccv/ablation/abl_all.out 2>&1 & +nohup python cvae.py --name="lhunor_abl_all_gnn" > nohup/iccv/ablation/abl_all_gnn.out 2>&1 & nohup python upscale_test.py --name="iccv_32_2" > nohup/iccv/upscale/iccv_32_2.out 2>&1 & nohup python upscale_test.py --name="iccv_64_2" > nohup/iccv/upscale/iccv_64_2.out 2>&1 & diff --git a/configs/iccv/config_abl_all.py b/configs/iccv/config_abl_all.py index 5c9a412..4e9f18d 100644 --- a/configs/iccv/config_abl_all.py +++ b/configs/iccv/config_abl_all.py @@ -19,10 +19,10 @@ 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'] = '3' -config['name'] = 'abl_all' +config['gpu'] = '2' +config['name'] = 'abl_all_gnn' config['notes'] = '' -config['tags'] = ['iccv', 'abl_all'] +config['tags'] = ['iccv', 'abl_all_gnn'] 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']: @@ -44,9 +44,9 @@ config['static_model'] = 'checkpoints/iccv/static_iccv_1' config['body_model'] = 'checkpoints/iccv/body_iccv_1' config['curvature_model'] = 'checkpoints/iccv/curv_iccv_1' -config['static_model'] = 'checkpoints/post/static_post_2' -config['body_model'] = 'checkpoints/post/body_post_2' -config['curvature_model'] = 'checkpoints/post/curv_post_1' +# config['static_model'] = 'checkpoints/post/static_post_2' +# config['body_model'] = 'checkpoints/post/body_post_2' +# config['curvature_model'] = 'checkpoints/post/curv_post_1' config['smpl_res'] = 14475 @@ -54,7 +54,7 @@ if config['is_cape']: config['smpl_res'] = 3638 -config['pretrained_model'] = '' +config['pretrained_model'] = 'checkpoints/iccv/abl_all' ############### config['lr'] = 0.0001# 0.0001 diff --git a/cvae.py b/cvae.py index 846e14d..d172c1a 100644 --- a/cvae.py +++ b/cvae.py @@ -47,7 +47,7 @@ from custom_utils.custom_utils import Accumulator def init_wandb(): - run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags']) + run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae.py') artifact.add_file('models/cvae_model.py') @@ -87,7 +87,7 @@ def train_step(batch): else: mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed'] - pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True) + pred, mean, logvar = 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=True) @@ -107,34 +107,29 @@ def train_step(batch): e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio) - l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2])) l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1)) n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) - # conc_loss = concept_loss(pred, batch['uv_unposed']) - kl = kl_loss(mean, logvar) - reconstr_loss = l2_loss + reconstr_loss = l2_mesh_loss if config['use_extra_losses']: if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True: - reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff'] + reconstr_loss = reconstr_loss + (e_loss) * config['reg_loss_coeff'] else: - reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff'] + reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff'] if config['use_norm_loss'] and step > 2000: reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff'] - loss = reconstr_loss + g_balancing_coeff * kl + loss = reconstr_loss - mean = tf.reduce_mean(mean) - logvar = tf.reduce_mean(logvar) # 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) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) - gradients = tape.gradient(loss, model.trainable_variables) - optimizer.apply_gradients(zip(gradients, model.trainable_variables)) + gradients = tape.gradient(loss, smoother.trainable_variables) + optimizer.apply_gradients(zip(gradients, smoother.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, - 'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss, + 'train/mean': mean, 'train/logvar': logvar, 'train/var': np.e ** logvar, + 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss, # 'train/ms_ssim': ms_ssim, } @@ -166,22 +161,16 @@ def valid_step(batch): e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio) - l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2])) l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1)) n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) - # conc_loss = concept_loss(pred, batch['uv_unposed']) - kl = kl_loss(mean, logvar) - reconstr_loss = l2_loss + reconstr_loss = l2_mesh_loss if config['use_extra_losses']: - reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff'] + reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff'] if config['use_norm_loss']: reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff'] - loss = reconstr_loss + g_balancing_coeff * kl - - mean = tf.reduce_mean(mean) - logvar = tf.reduce_mean(logvar) + loss = reconstr_loss # 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) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) @@ -189,7 +178,7 @@ def valid_step(batch): return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric, 'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/collision_ratio': c_ratio, - 'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar, + 'valid/mean': mean, 'valid/logvar': logvar, 'valid/var': np.e ** logvar, 'valid/normal_loss': n_loss,# 'valid/ms_ssim': ms_ssim, } @@ -291,7 +280,6 @@ def plot_everything(batch: dict, type: str): gt_posed = [batch['mesh_posed'].numpy()[i, mask[i], :] for i in range(mask.shape[0])] - body = batch['mesh_body_posed'].numpy() mesh_posed_fixed = fix_collision(mesh_posed, body, mask).numpy() @@ -525,6 +513,7 @@ if __name__ == '__main__': is_cape=config['is_cape']) model = CVAE(config) + model.trainable = False smoother = Smoother() optimizer = tf.optimizers.Adam(learning_rate=config['lr']) @@ -558,17 +547,6 @@ if __name__ == '__main__': 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): - # 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'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017): if balancing_activated_step == 0: balancing_activated_step = step diff --git a/models/gnn.py b/models/gnn.py index ad9d495..07d65c7 100644 --- a/models/gnn.py +++ b/models/gnn.py @@ -16,12 +16,11 @@ class Smoother(tf.keras.Model): def __init__(self, ): super().__init__() self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()) - self.conv1 = GeneralConv(channels=32, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True) - self.conv2 = GeneralConv(channels=128, batch_norm=True, dropout=0.1, aggregate='mean', activation='prelu', use_bias=True) - self.conv3 = GeneralConv(channels=64, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True) - self.conv4 = GeneralConv(channels=16, batch_norm=True, dropout=0.0, aggregate='mean', activation='prelu', use_bias=True) - - self.conv5 = GCNConv(3, activation='tanh', use_bias=False) + self.conv1 = GCNConv(16, activation='tanh', use_bias=False) + self.conv2 = GCNConv(16, activation='tanh', use_bias=False) + self.conv3 = GCNConv(16, activation='tanh', use_bias=False) + # self.conv4 = GCNConv(16, activation='tanh', use_bias=False) + # self.conv5 = GCNConv(16, activation='tanh', use_bias=False) self.conv6 = GCNConv(3, activation='tanh', use_bias=False) def call(self, inputs, training=None): @@ -29,8 +28,8 @@ class Smoother(tf.keras.Model): 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.conv4([x, self.a], training=training) + # x = self.conv5([x, self.a], training=training) x = self.conv6([x, self.a], training=training) return x