loss 分支train

brep2sdf/train.py

@@ -281,6 +281,8 @@ class Trainer:
                     #if check_tensor(normals, "Normals (Loss Input)", epoch, step): raise ValueError("Bad normals before loss")
                     #if check_tensor(points, "Points (Loss Input)", epoch, step): raise ValueError("Bad points before loss")
                     logger.gpu_memory_stats("计算损失前")
+
+                    if args.only_zero_surface:
                         loss, loss_details = self.loss_manager.compute_loss(
                             mnfld_pnts,
                             nonmnfld_pnts,
@@ -290,6 +292,13 @@ class Trainer:
                             nonmnfld_pred,
                             psdf
                         )       
+                    else:
+                        loss, loss_details = self.loss_manager.compute_loss_stage1(
+                            mnfld_pnts,
+                            normals, # 传递检查过的 normals
+                            gt_sdf,
+                            mnfld_pred
+                        )              
                 else:
                     loss = torch.nn.functional.mse_loss(pred_sdf, gt_sdf)
 
@@ -696,6 +705,7 @@ class Trainer:
                     #if check_tensor(normals, "Normals (Loss Input)", epoch, step): raise ValueError("Bad normals before loss")
                     #if check_tensor(points, "Points (Loss Input)", epoch, step): raise ValueError("Bad points before loss")
                     #logger.gpu_memory_stats("计算损失前")
+                    if args.only_zero_surface:
                         loss, loss_details = self.loss_manager.compute_loss(
                             mnfld_pnts,
                             nonmnfld_pnts,
@@ -705,6 +715,13 @@ class Trainer:
                             nonmnfld_pred,
                             psdf
                         )       
+                    else:
+                        loss, loss_details = self.loss_manager.compute_loss_stage1(
+                            mnfld_pnts,
+                            normals, # 传递检查过的 normals
+                            gt_sdf,
+                            mnfld_pred
+                        )               
                     #logger.gpu_memory_stats("计算损失后")
                 else:
                     loss = torch.nn.functional.mse_loss(pred_sdf, gt_sdf)