From 65fd3d9dedccf4a6bcb6d7324745052e3e7628bc Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 1 Mar 2023 08:45:33 +0000 Subject: [PATCH] training fix --- cvae.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cvae.py b/cvae.py index 5315bbd..cdbb8af 100644 --- a/cvae.py +++ b/cvae.py @@ -128,8 +128,8 @@ def train_step(batch): # 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)) + gradients = tape.gradient(loss, model.trainable_variables + smoother.trainable_variables) + optimizer.apply_gradients(zip(gradients, model.trainable_variables + smoother.trainable_variables)) return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric, 'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,