configs save

This commit is contained in:
2023-07-02 01:06:14 +02:00
parent 05fc297808
commit d66e6d5a77
9 changed files with 555 additions and 12 deletions
+1 -1
View File
@@ -94,7 +94,7 @@ def train_step(batch):
mesh_orig = batch['mesh_static']
l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1))
e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }