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

large refactor

Browse Source
This commit is contained in:
luni
2022-09-25 08:31:13 +00:00
parent 84ec5a056e
commit 92891f5b26
33 changed files with 1661 additions and 1635 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+38 -138
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -31,83 +31,37 @@ from load_models import load_generator, load_discriminator
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAE
from configs.config_cvae import config
name='cvae_17'
def init_wandb():
run = wandb.init(project='vae', name=name,
notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \
run = wandb.init(project='vae', name=config['name'],
notes='\
offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
retrained body encoder with deoriented bodies \
use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
fixed gender issue and retrained body encoder, FINAL\
remove shape from unposed garment\
training on 100 garments with offsets over the average of these')
training on 100 garments with offsets over the average of these\
fixed autoencoders!!!!!!!!!')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('models/cvae_model.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.py')
artifact.add_file('configs/config.py')
artifact.add_file('configs/config_cvae.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('dataset.py')
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,2] += 1
collision[:,2] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,2],
z=mesh_gt[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,2],
z=mesh_pred_unposed[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,2],
z=collision[:,1],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
def log_normal_pdf(sample, mean, logvar, raxis=1):
log2pi = tf.math.log(2. * np.pi)
@@ -122,13 +76,14 @@ def kl_loss(mean, logvar):
return kl_loss
@tf.function
# @tf.function
def train_step(batch):
with tf.GradientTape() as tape:
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']], training=True, use_mask=True)
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']],
training=True, use_mask=True)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
@@ -141,38 +96,27 @@ def train_step(batch):
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']))
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])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
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, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
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))
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])
tv_loss = tf.image.total_variation(pred)
reconstr_loss = l2_loss
reconstr_loss = tf.reduce_mean(reconstr_loss)
kl = kl_loss(mean, logvar)
loss = reconstr_loss + balancing_coeff * kl
# 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)
fft_loss = tf.reduce_mean(fft_loss)
tv_loss = tf.reduce_mean(tv_loss)
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)
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 = 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'])
euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
gradients = tape.gradient(loss, model.trainable_variables)
# gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
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,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2,
'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
'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 }
@@ -183,7 +127,7 @@ 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']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
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']))
@@ -197,10 +141,6 @@ def valid_step(batch):
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
# logpx_z = -reconstr_loss
# logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
# logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
# loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
kl = kl_loss(mean, logvar)
# loss = reconstr_loss / (2 * (gamma ** 2)) + kl
loss = reconstr_loss + balancing_coeff * kl
@@ -224,9 +164,9 @@ def plot_everything(batch: dict, type: str):
mesh_pred = uv2mesh(pred).numpy()
mesh_orig = batch['mesh_unposed']
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
orig_img = batch['uv_unposed']
uv_static_img = batch['uv_static']
uv_body_posed_img = batch['uv_body_posed']
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)]
@@ -259,8 +199,6 @@ def plot_everything(batch: dict, type: str):
return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def sample(folder):
# run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
load_model(model,folder)
@@ -274,7 +212,7 @@ def sample(folder):
pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
mesh_pred = uv2mesh(pred)
mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
mesh_orig = batch['mesh_posed'][0].numpy()
mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
@@ -286,9 +224,8 @@ def sample(folder):
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
pred_zero = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2]))
mesh_pred_zero = uv2mesh(pred_zero)
mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
@@ -382,67 +319,42 @@ def sample(folder):
print()
if __name__ == '__main__':
init_wandb()
# init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
batch_size = 4
batch_size = config['batch_size']
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
'uv_curvature': (512, 256),
}
signature_static = { 'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=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']])
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'])
model = CVAE(config)
model = CVAE(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
latent_dim=128, static_model='checkpoints/ae_static_4',
body_model='checpoints/ae_body_8',
curvature_model='checpoints/ae_curvature_7')
# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
optimizer = tf.optimizers.Adam(learning_rate=0.0001)
balancing_coeff = 0.0001
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
balancing_coeff = config['starting_balancing_coeff']
reg_coeff = 1
folder = 'models/ae_logs/'
valid_iterator = iter(dataset_valid)
gamma = 1
balancing_activated_step = 0
# avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy')
# sample('checkpoints/cvae_13')
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
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, signature_dynamic))
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
train_loss = train_step(batch)
@@ -464,19 +376,7 @@ if __name__ == '__main__':
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.015:
# reg_coeff = reg_coeff + 100 / 15000
# if step > 15000 and balancing_coeff <= 1e-3:
# balancing_coeff = 1e-2
# if step > 20000 and balancing_coeff <= 1e-3:
# balancing_coeff = 1e-2
# if step > 25000 and balancing_coeff <= 1e-2:
# balancing_coeff = 1e-1
# if step > 30000 and balancing_coeff <= 1e-1:
# balancing_coeff = 1
# gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss']
balancing_coeff = balancing_coeff + 1 / 15000
to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }