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

fix training

Browse Source
This commit is contained in:
luni
2023-02-22 22:51:00 +00:00
parent c5d576f11f
commit 21bd18cbe2
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+2 -2
View File
@@ -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))) # 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']) 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) gradients = tape.gradient(loss, model.trainable_variables + smoother.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.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, 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/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,