From 01f605f7b78da830c7998d46d7613eeff16b2635 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 12 Oct 2022 21:07:19 +0000 Subject: [PATCH] update cvae --- cvae.py | 200 ++++++++++++++++++++++++++++++++++---------------------- 1 file changed, 123 insertions(+), 77 deletions(-) diff --git a/cvae.py b/cvae.py index 228bbf5..74527f0 100644 --- a/cvae.py +++ b/cvae.py @@ -21,7 +21,7 @@ import wandb import random import numpy as np import scipy as sp -from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces +from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces import plotly.graph_objects as go import plotly.express as ex from PIL import Image @@ -44,7 +44,7 @@ from models.cvae_model import CVAE def init_wandb(): - run = wandb.init(project='vae', name=config['name'], notes=config['notes']) + run = wandb.init(project='cvae', name=config['name'], notes=config['notes']) artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae.py') artifact.add_file('models/cvae_model.py') @@ -85,71 +85,66 @@ def train_step(batch): shape = batch['pose_shape'][:, 72:] gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) - e_loss = edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']) - c_loss = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) - l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2]) - l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2]) - # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2) - ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)) - # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2)) + # is the output already masked? + + e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])) + l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) + kl = kl_loss(mean, logvar) reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'] - reconstr_loss = tf.reduce_mean(reconstr_loss) - kl = kl_loss(mean, logvar) - loss = reconstr_loss + balancing_coeff * kl - ms_ssim = tf.reduce_mean(ms_ssim) - l1_loss = tf.reduce_mean(l1_loss) - l2_loss = tf.reduce_mean(l2_loss) - e_loss = tf.reduce_mean(e_loss) - c_loss = tf.reduce_mean(c_loss) + loss = reconstr_loss + g_balancing_coeff * kl + 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)) - return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, + return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, 'train/reconstr_loss': reconstr_loss, - 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar, - 'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss } + 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar, + 'train/l2_loss': l2_loss } @tf.function def valid_step(batch): - uv_curvature = tf.one_hot(batch['uv_curvature'], 7) - pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']]) + pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) mesh_pred = uv2mesh(pred) mesh_orig = batch['mesh_unposed'] - e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) - l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2]) - b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask']) - # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2) - ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3) - fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3]) - # mesh_orig_2 = uv2mesh(batch['uv_unposed']) - 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']) - metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) - - reconstr_loss = l1_loss + + shape = batch['pose_shape'][:, 72:] + gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) + shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) + + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] + mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) + + # is the output already masked? + + e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])) + l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2])) kl = kl_loss(mean, logvar) - # loss = reconstr_loss / (2 * (gamma ** 2)) + kl - loss = reconstr_loss + balancing_coeff * kl - ms_ssim = tf.reduce_mean(ms_ssim) - l1_loss = tf.reduce_mean(l1_loss) - fft_loss = tf.reduce_mean(fft_loss) - b_loss = tf.reduce_mean(b_loss) - reconstr_loss = tf.reduce_mean(reconstr_loss) + + reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'] + loss = reconstr_loss + g_balancing_coeff * kl + mean = tf.reduce_mean(mean) logvar = tf.reduce_mean(logvar) - return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, - 'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_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']) + + return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, + 'train/reconstr_loss': reconstr_loss, 'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar } @@ -315,6 +310,57 @@ def sample(folder): print() + + + +def test(): + load_model(model, 'checkpoints/' + config['name'].split('_')[0]) + + dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False) + static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False) + + step = 0 + losses = [] + + for batch in dataset_test: + batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) + + test_loss = valid_step(batch) + losses.append(test_loss) + + to_log = { **{ 'current_step': step }, **test_loss } + if step % 100 == 0: + figs = plot_everything(batch, 'test') + to_log = { **to_log, **figs } + + wandb.log(to_log, step = step) + + if step > 30: + break + + print(step) + step += 1 + + # averages = dict() + path_to_plotly_html = "temp/plotly_figure.html" + + histograms = dict() + for key in test_loss.keys(): + accumulated = [losses[i][key].numpy() for i in range(len(losses))] + fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean())) + fig.write_html(path_to_plotly_html, auto_play = False) + + histograms['all/' + key.split('/')[-1] + '_hist'] = wandb.Html(path_to_plotly_html) + # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean() + + + + wandb.log({ **histograms }, step = step) + wandb.finish() + pass + + + if __name__ == '__main__': init_wandb() allow_memory_growth() @@ -327,55 +373,55 @@ if __name__ == '__main__': factory = DatasetFactory() - dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']]) - static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']]) - dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic']) - static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static']) + dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) + static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) + dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) + static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) model = CVAE(config) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) - balancing_coeff = config['starting_balancing_coeff'] - reg_coeff = 1 - valid_iterator = iter(dataset_valid) - gamma = 1 balancing_activated_step = 0 + g_balancing_coeff = config['starting_balancing_coeff'] # sample('checkpoints/cvae_13') + test() - ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3) + # ckpt = tf.train.Checkpoint(model=model) + # manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3) - step = 0 - for batch in dataset_train: - batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) + # valid_iterator = iter(dataset_valid) + # step = 0 + # for batch in dataset_train: + # batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) - train_loss = train_step(batch) + # train_loss = train_step(batch) - # data_valid = valid_iterator.get_next() - # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) - # valid_loss = valid_step(data_valid) - valid_loss = {} + # 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) - to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } - if step % 2500 == 0: - train_figs = plot_everything(batch, 'train') - # valid_figs = plot_everything(data_valid, 'valid') - valid_figs = {} - to_log = { **to_log, **train_figs, **valid_figs } + # 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 % 10000 == 0: - manager.save(checkpoint_number=step) + # if step % 10000 == 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.01: - if balancing_activated_step == 0: - balancing_activated_step = step - balancing_coeff = balancing_coeff + 1 / 15000 + # if train_loss['train/euclid_error'] < 0.01: + # if balancing_activated_step == 0: + # balancing_activated_step = step + # g_balancing_coeff = g_balancing_coeff + 1 / 15000 - to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } + # to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } } - wandb.log(to_log, step = step) + # wandb.log(to_log, step = step) - print(step) + # print(step) - step += 1 \ No newline at end of file + # step += 1 + \ No newline at end of file