update pyramid
This commit is contained in:
+82
-57
@@ -43,7 +43,7 @@ from models.pyramid_model import Pyramid
|
||||
|
||||
|
||||
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_pyramid.py')
|
||||
@@ -111,6 +111,36 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
|
||||
),
|
||||
row=1, col=1+len(mesh_pred_unposed_levels)+1+1)
|
||||
|
||||
if len(mesh_pred_unposed_levels) == 1:
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
|
||||
)
|
||||
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
fig.update_layout(scene = helper,
|
||||
scene2 = helper,
|
||||
scene3 = helper,
|
||||
scene4 = helper,
|
||||
)
|
||||
if len(mesh_pred_unposed_levels) == 2:
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
|
||||
)
|
||||
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
fig.update_layout(scene = helper,
|
||||
scene2 = helper,
|
||||
scene3 = helper,
|
||||
scene4 = helper,
|
||||
scene5 = helper,
|
||||
)
|
||||
if len(mesh_pred_unposed_levels) == 3:
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
@@ -129,22 +159,6 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
|
||||
scene5 = helper,
|
||||
scene6 = helper,
|
||||
)
|
||||
if len(mesh_pred_unposed_levels) == 2:
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
|
||||
)
|
||||
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
fig.update_layout(scene = helper,
|
||||
scene2 = helper,
|
||||
scene3 = helper,
|
||||
scene4 = helper,
|
||||
scene5 = helper,
|
||||
)
|
||||
if len(mesh_pred_unposed_levels) == 4:
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
@@ -230,8 +244,8 @@ 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']))
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
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)))
|
||||
#ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
|
||||
# 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)))
|
||||
ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
|
||||
kl = kl_loss(mean, logvar)
|
||||
|
||||
reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
|
||||
@@ -257,41 +271,54 @@ def train_step(batch):
|
||||
|
||||
@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']])
|
||||
levels_input, conds = group_inputs_for_levels(batch)
|
||||
out_levels = model([levels_input, conds], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
|
||||
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_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))
|
||||
output = dict()
|
||||
reconstr_loss_total = 0
|
||||
kl_total = 0
|
||||
|
||||
#for i, res in enumerate(config['pyramid']['resolutions']):
|
||||
i = -1
|
||||
res = config['pyramid']['resolutions'][-1]
|
||||
res_modifier = '_' + str(res) if res != 512 else ''
|
||||
pred, mean, logvar = out_levels[i]
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
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]))
|
||||
# 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)))
|
||||
ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
|
||||
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_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
|
||||
# reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
|
||||
kl_total += 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,
|
||||
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
|
||||
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'])
|
||||
|
||||
output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric,
|
||||
'valid/ms_ssim' + res_modifier: ms_ssim, 'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
|
||||
'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss,
|
||||
'valid/euclid_error': metric } }
|
||||
|
||||
loss = reconstr_loss_total + balancing_coeff * kl_total
|
||||
|
||||
return { **{ 'valid/loss': loss, 'valid/reconstr_loss_total': reconstr_loss_total, 'valid/kl_total': kl_total }, **output }
|
||||
|
||||
|
||||
def plot_everything(batch: dict, type: str):
|
||||
@@ -484,7 +511,7 @@ if __name__ == '__main__':
|
||||
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'])
|
||||
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions)
|
||||
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'])
|
||||
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
|
||||
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions)
|
||||
|
||||
model: Pyramid = Pyramid(config)
|
||||
|
||||
@@ -511,12 +538,12 @@ if __name__ == '__main__':
|
||||
# model.levels[0].upscale_enhancer.trainable = False
|
||||
|
||||
# the last level doesn't have an upscale enhancer as it's enhanced for the final resolutin directly
|
||||
# might now be necessary, as enhancing to final or a bit higher resolution seems to be the same
|
||||
# might not be necessary, as enhancing to final or a bit higher resolution seems to be the same
|
||||
for i in range(len(model.levels)-1):
|
||||
load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
|
||||
model.levels[i].upscale_enhancer.trainable = False
|
||||
# load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
# model.upscale_enhancer.trainable = False
|
||||
load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
model.upscale_enhancer.trainable = False
|
||||
|
||||
for i in range(len(model.levels)-1):
|
||||
model.levels[i].trainable = False
|
||||
@@ -532,16 +559,14 @@ if __name__ == '__main__':
|
||||
|
||||
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 % 5000 == 0 or (step < 5000 and step % 500 == 0):
|
||||
train_figs = plot_everything(batch, 'train')
|
||||
# valid_figs = plot_everything(data_valid, 'valid')
|
||||
valid_figs = {}
|
||||
valid_figs = plot_everything(data_valid, 'valid')
|
||||
to_log = { **to_log, **train_figs, **valid_figs }
|
||||
|
||||
if step % 10000 == 0:
|
||||
|
||||
Reference in New Issue
Block a user