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

working pyramid pipeline

Browse Source
This commit is contained in:
luni
2022-09-27 21:26:15 +00:00
parent 1da83294ee
commit fa8f24f0c5
10 changed files with 997 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_pyramid_old.py
+520
View File
@@ -0,0 +1,520 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
# Currently, memory growth needs to be the same across GPUs
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
except RuntimeError as e:
# Memory growth must be set before GPUs have been initialized
print(e)
import time
import argparse
import wandb
import random
import numpy as np
import scipy as sp
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
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 dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_pyramid_model import Pyramid
name='cvae_pyr_4_d32'
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. \
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\
remove shape from unposed garment\
training on 100 garments with offsets over the average of these\
l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this\
low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input\
not using filled in versions, but using custom MLP upscaling')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae_pyramid.py')
artifact.add_file('models/cvae_model.py')
artifact.add_file('models/cvae_pyramid_model.py')
artifact.add_file('models/upscaler.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.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, mesh_pred_unposed_low, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,2] += 1
mesh_pred_unposed_low[:,2] += 2
collision[:,2] += 3
fig = go.Figure(data=[
go.Mesh3d(
x=mesh_gt[:,0],
y=mesh_gt[:,2],
z=mesh_gt[:,1],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
x=mesh_pred_unposed[:,0],
y=mesh_pred_unposed[:,2],
z=mesh_pred_unposed[:,1],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
x=mesh_pred_unposed_low[:,0],
y=mesh_pred_unposed_low[:,2],
z=mesh_pred_unposed_low[:,1],
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
x=collision[:,0],
y=collision[:,2],
z=collision[:,1],
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=3, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,5],),
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)
res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
return res
def kl_loss(mean, logvar):
kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
kl_loss = tf.reduce_mean(kl_loss)
kl_loss *= -0.5
return kl_loss
# @tf.function
def train_step(batch):
with tf.GradientTape() as tape:
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7)
pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'],
batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier],])
mesh_pred_full = uv2mesh(pred_full)
mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full)
mesh_pred_low = uv2mesh(pred_low)
mesh_pred_posed_low = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_low)
mesh_orig = batch['mesh_unposed']
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_posed_full = mesh_pred_posed_full + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred_posed_low = mesh_pred_posed_low + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
e_loss = tf.reduce_mean(edge_loss(mesh_pred_full, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed_full, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred_full)), axis=-1), axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full) ** 2), axis=-1), axis=[1, 2])
ms_ssim = tf.image.ssim_multiscale(pred_full + 1, batch['uv_unposed'] + 1, max_val=3)
kl = kl_loss(mean_full, logvar_full)
pred_low = pred_low * batch['uv_mask' + resolution_modifier][:,:,:,tf.newaxis]
e_loss_low = tf.reduce_mean(edge_loss(mesh_pred_low, mesh_orig, batch['mesh_mask']))
c_loss_low = tf.reduce_mean(collision_penalty(mesh_pred_posed_low, batch['mesh_body_posed'], batch['mesh_mask']))
l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + resolution_modifier] - pred_low)), axis=-1), axis=[1, 2])
l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + resolution_modifier] - pred_low) ** 2), axis=-1), axis=[1, 2])
ms_ssim_low = tf.image.ssim(pred_low + 1, batch['uv_unposed' + resolution_modifier] + 1, max_val=3)
kl_low = kl_loss(mean_full, logvar_full)
reconstr_loss = l2_loss + l2_loss_low
reconstr_loss = tf.reduce_mean(reconstr_loss)
loss = reconstr_loss + balancing_coeff * (kl + kl_low)
ms_ssim = tf.reduce_mean(ms_ssim)
l1_loss = tf.reduce_mean(l1_loss)
l2_loss = tf.reduce_mean(l2_loss)
mean = tf.reduce_mean(mean_full)
logvar = tf.reduce_mean(logvar_full)
ms_ssim_low = tf.reduce_mean(ms_ssim_low)
l1_loss_low = tf.reduce_mean(l1_loss_low)
l2_loss_low = tf.reduce_mean(l2_loss_low)
mean_low = tf.reduce_mean(mean_low)
logvar_low = tf.reduce_mean(logvar_low)
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_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_low = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_low) ** 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_full) ** 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,
'train/euclid_error_reconst': metric_2, 'train/reconstr_loss': reconstr_loss, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl,
'train/var': np.e ** logvar, 'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss,
'train/edge_loss' + resolution_modifier: e_loss_low, 'train/collision_loss' + resolution_modifier: c_loss_low, 'train/l1_loss' + resolution_modifier: l1_loss_low, 'train/euclid_error' + resolution_modifier: metric_low, 'train/ms_ssim' + resolution_modifier: ms_ssim_low,
'train/mean' + resolution_modifier: mean_low, 'train/logvar' + resolution_modifier: logvar_low, 'train/kl_loss' + resolution_modifier: kl_low,
'train/var' + resolution_modifier: np.e ** logvar_low, 'train/l2_loss' + resolution_modifier: l2_loss_low }
@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']])
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
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
# 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
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)
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,
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
def plot_everything(batch: dict, type: str):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7)
pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'],
batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier]])
mesh_pred_full = uv2mesh(pred_full)
mesh_pred_low = uv2mesh(pred_low)
mesh_orig = batch['mesh_unposed']
orig_img = batch['uv_unposed']
uv_static_img = batch['uv_static']
uv_body_posed_img = batch['uv_body_posed']
pred_img = pred_full
pred_low_up = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
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])])
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred_full.numpy(), F, mask)
garment_unposed_low, faces = maskFacesBatch(mesh_pred_low.numpy(), F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
garment_unposed_low = [garment_unposed_low[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
for i in range(mask.shape[0]):
mesh_pred_full[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
body = batch['mesh_body_posed'].numpy()
mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy()
garment_posed = [mesh_posed_full[i, mask[i], :] for i in range(body.shape[0])]
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], garment_unposed_low[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
mean_hist = wandb.Histogram(sequence=mean_full.numpy().mean(axis=0))
logvar_hist = wandb.Histogram(sequence=logvar_full.numpy().mean(axis=0))
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)
count = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
nr_of_samples=2
pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
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_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]
body = batch['mesh_body_posed'][0].numpy()
body_color = np.ones_like(body, dtype=np.int32)*100
F = np.load("custom_utils/smpl_faces.npy")
mask = batch['mesh_mask'].numpy().astype(bool)[0]
_, faces = _maskFaces(np.ones((14475, 3)), F, mask)
pred_zero = model.sample(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_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
meshes = []
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=mesh_orig_reposed[mask,0],
y=mesh_orig_reposed[mask,2],
z=mesh_orig_reposed[mask,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
for i in range(nr_of_samples):
x = mesh_posed[i][mask, :].copy()
# x[:, 2] += 1 * (i + 1)
x[:, 2] += 1
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,2],
z=x[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
for i in range(nr_of_samples):
x = mesh_posed_zero[i][mask, :].copy()
# x[:, 2] += 1 * (i + 1)
x[:, 2] += 2
meshes.append(go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,2],
z=x[:,1],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
))
# for i in range(5):
# x = body.copy()
# x[:, 2] += 1 * i
# meshes.append(go.Mesh3d(
# # 8 vertices of a cube
# x=x[:,0],
# y=x[:,2],
# z=x[:,1],
# # i, j and k give the vertices of triangles
# i = F[:,0],
# j = F[:,1],
# k = F[:,2],
# vertexcolor=body_color
# ))
fig = go.Figure(data=meshes)
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
# fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
# yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
# wandb.log({'plot_' + str(count): fig})
fig.show()
pred = pred.numpy()
frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
save_all=True, duration=500, loop=0)
if count > 10:
break
print(count)
count += 1
print()
if __name__ == '__main__':
# init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
batch_size = 4
low_resolution = 32
resolution_modifier = '_' + str(low_resolution)
smpl = SMPL()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (512, 256, 3),
'uv_body_posed': (512, 256, 3),
'uv_body_posed' + resolution_modifier: (low_resolution, low_resolution//2, 3),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (512, 256, 3),
'uv_unposed' + resolution_modifier: (low_resolution, low_resolution//2, 3),
'uv_curvature': (512, 256),
'uv_curvature' + resolution_modifier: (low_resolution, low_resolution//2),
}
signature_static = { 'uv_static': (512, 256, 3),
'uv_static' + resolution_modifier: (low_resolution, low_resolution//2, 3),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
'uv_mask' + resolution_modifier: (low_resolution, low_resolution//2,),
# 'uv_shaped_body': (512, 256, 3),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, low_resolution])
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_resolution])
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)
model = Pyramid(latent_dim=128,
low_res=low_resolution,
freeze_encoders=True,
static_model='checkpoints/ae_static_6',
body_model='checkpoints/ae_body_8',
curvature_model='checkpoints/ae_curvature_7',
upscaler_model='checkpoints/up_2_3')
# 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
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)
step = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
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 = {}
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 }
if step % 10000 == 0:
manager.save(checkpoint_number=step)
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']
to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
wandb.log(to_log, step = step)
print(step)
step += 1