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

training fix

Browse Source
This commit is contained in:
luni
2023-03-01 08:45:33 +00:00
parent 53a9382e1c
commit 65fd3d9ded
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)))
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,