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

update cvae

Browse Source
This commit is contained in:
luni
2022-10-12 21:07:19 +00:00
parent a9e8072f90
commit 01f605f7b7
1 changed files with 123 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+123 -77
View File
@@ -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
# step += 1