Refactor validation logging and loss normalization (#1126)

Author: YonghaoZhao722

micro_sam/training/joint_sam_trainer.py

@@ -136,6 +136,7 @@ def _validate_impl(self, forward_context):
 
         val_iteration = 0
         metric_val, loss_val, model_iou_val = 0.0, 0.0, 0.0
+        mask_loss_val, iou_loss_val, unetr_loss_val = 0.0, 0.0, 0.0
 
         with torch.no_grad():
             for x, y in self.val_loader:
@@ -155,17 +156,23 @@ def _validate_impl(self, forward_context):
 
                 loss_val += loss.item()
                 metric_val += metric.item() + (unetr_metric.item() / 3)
+                mask_loss_val += mask_loss.item()
+                iou_loss_val += iou_regression_loss.item()
                 model_iou_val += model_iou.item()
+                unetr_loss_val += unetr_loss.item()
                 val_iteration += 1
 
         loss_val /= len(self.val_loader)
         metric_val /= len(self.val_loader)
+        mask_loss_val /= len(self.val_loader)
+        iou_loss_val /= len(self.val_loader)
         model_iou_val /= len(self.val_loader)
+        unetr_loss_val /= len(self.val_loader)
 
         if self.logger is not None:
             self.logger.log_validation(
                 self._iteration, metric_val, loss_val, x, labels_instances, sampled_binary_y,
-                mask_loss, iou_regression_loss, model_iou_val, unetr_loss
+                mask_loss_val, iou_loss_val, model_iou_val, unetr_loss_val
             )
 
         return metric_val

micro_sam/training/sam_trainer.py

@@ -134,6 +134,9 @@ def _compute_loss(self, batched_outputs, y_one_hot):
         - The IOU loss: L2 loss between the predicted IOU and the actual IOU of prediction and target.
         """
         mask_loss, iou_regression_loss = 0.0, 0.0
+        batch_size = len(batched_outputs)
+        if batch_size == 0:
+            raise RuntimeError("Got empty batch outputs in loss computation.")
 
         # Loop over the batch.
         for batch_output, targets in zip(batched_outputs, y_one_hot):
@@ -163,6 +166,9 @@ def _compute_loss(self, batched_outputs, y_one_hot):
             mask_loss = mask_loss + torch.mean(dice_scores)
             iou_regression_loss = iou_regression_loss + iou_score
 
+        # Normalize by batch size so that loss/metric are comparable across batch sizes.
+        mask_loss = mask_loss / batch_size
+        iou_regression_loss = iou_regression_loss / batch_size
         loss = mask_loss + iou_regression_loss
 
         return loss, mask_loss, iou_regression_loss
@@ -448,6 +454,7 @@ def _validate_impl(self, forward_context):
 
         val_iteration = 0
         metric_val, loss_val, model_iou_val = 0.0, 0.0, 0.0
+        mask_loss_val, iou_loss_val = 0.0, 0.0
 
         with torch.no_grad():
             for x, y in self.val_loader:
@@ -459,19 +466,23 @@ def _validate_impl(self, forward_context):
 
                 loss_val += loss.item()
                 metric_val += metric.item()
+                mask_loss_val += mask_loss.item()
+                iou_loss_val += iou_regression_loss.item()
                 model_iou_val += model_iou.item()
                 val_iteration += 1
 
         loss_val /= len(self.val_loader)
         metric_val /= len(self.val_loader)
+        mask_loss_val /= len(self.val_loader)
+        iou_loss_val /= len(self.val_loader)
         model_iou_val /= len(self.val_loader)
         print()
         print(f"The Average Dice Score for the Current Epoch is {1 - metric_val}")
 
         if self.logger is not None:
             self.logger.log_validation(
                 self._iteration, metric_val, loss_val, x, y,
-                sampled_binary_y, mask_loss, iou_regression_loss, model_iou_val
+                sampled_binary_y, mask_loss_val, iou_loss_val, model_iou_val
             )
 
         return metric_val