Correct masked mean computation in AUCMLoss and update docstrings

Signed-off-by: Shubham Chandravanshi <[email protected]>

Author: shubham-61969

monai/losses/aucm_loss.py

@@ -11,7 +11,6 @@
 
 from __future__ import annotations
 
-
 import torch
 import torch.nn as nn
 from torch.nn.modules.loss import _Loss
@@ -93,6 +92,9 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             input: the shape should be B1HW[D], where the channel dimension is 1 for binary classification.
             target: the shape should be B1HW[D], with values 0 or 1.
 
+        Returns:
+            torch.Tensor: scalar AUCM loss.
+
         Raises:
             ValueError: When input or target have incorrect shapes.
         """
@@ -112,25 +114,29 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         if self.version == "v1":
             p = self.imratio if self.imratio is not None else pos_mask.mean()
             loss = (
-                (1 - p) * self._safe_mean((input - self.a) ** 2 * pos_mask)
-                + p * self._safe_mean((input - self.b) ** 2 * neg_mask)
+                (1 - p) * self._safe_mean((input - self.a) ** 2, pos_mask)
+                + p * self._safe_mean((input - self.b) ** 2, neg_mask)
                 + 2
                 * self.alpha
-                * (p * (1 - p) * self.margin + self._safe_mean(p * input * neg_mask - (1 - p) * input * pos_mask))
+                * (
+                    p * (1 - p) * self.margin
+                    + self._safe_mean(p * input * neg_mask - (1 - p) * input * pos_mask, pos_mask + neg_mask)
+                )
                 - p * (1 - p) * self.alpha**2
             )
         else:
             loss = (
-                self._safe_mean((input - self.a) ** 2 * pos_mask)
-                + self._safe_mean((input - self.b) ** 2 * neg_mask)
-                + 2 * self.alpha * (self.margin + self._safe_mean(input * neg_mask) - self._safe_mean(input * pos_mask))
+                self._safe_mean((input - self.a) ** 2, pos_mask)
+                + self._safe_mean((input - self.b) ** 2, neg_mask)
+                + 2 * self.alpha * (self.margin + self._safe_mean(input, neg_mask) - self._safe_mean(input, pos_mask))
                 - self.alpha**2
             )
 
         return loss
 
-    def _safe_mean(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Compute mean safely, returning 0 if tensor is empty."""
-        if tensor.numel() == 0 or tensor.count_nonzero() == 0:
+    def _safe_mean(self, tensor: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """Compute mean safely over masked elements."""
+        denom = mask.sum()
+        if denom == 0:
             return torch.tensor(0.0, device=tensor.device, dtype=tensor.dtype)
-        return tensor.sum() / tensor.count_nonzero()
+        return (tensor * mask).sum() / denom

tests/losses/test_aucm_loss.py

@@ -20,7 +20,10 @@
 
 
 class TestAUCMLoss(unittest.TestCase):
+    """Test cases for AUCMLoss."""
+
     def test_v1(self):
+        """Test AUCMLoss with version 'v1'."""
         loss_fn = AUCMLoss(version="v1")
         input = torch.randn(32, 1, requires_grad=True)
         target = torch.randint(0, 2, (32, 1)).float()
@@ -29,6 +32,7 @@ def test_v1(self):
         self.assertEqual(loss.ndim, 0)
 
     def test_v2(self):
+        """Test AUCMLoss with version 'v2'."""
         loss_fn = AUCMLoss(version="v2")
         input = torch.randn(32, 1, requires_grad=True)
         target = torch.randint(0, 2, (32, 1)).float()
@@ -37,31 +41,36 @@ def test_v2(self):
         self.assertEqual(loss.ndim, 0)
 
     def test_invalid_version(self):
+        """Test that invalid version raises ValueError."""
         with self.assertRaises(ValueError):
             AUCMLoss(version="invalid")
 
     def test_invalid_input_shape(self):
+        """Test that invalid input shape raises ValueError."""
         loss_fn = AUCMLoss()
         input = torch.randn(32, 2)  # Wrong channel
         target = torch.randint(0, 2, (32, 1)).float()
         with self.assertRaises(ValueError):
             loss_fn(input, target)
 
     def test_invalid_target_shape(self):
+        """Test that invalid target shape raises ValueError."""
         loss_fn = AUCMLoss()
         input = torch.randn(32, 1)
         target = torch.randint(0, 2, (32, 2)).float()  # Wrong channel
         with self.assertRaises(ValueError):
             loss_fn(input, target)
 
     def test_shape_mismatch(self):
+        """Test that mismatched shapes raise ValueError."""
         loss_fn = AUCMLoss()
         input = torch.randn(32, 1)
         target = torch.randint(0, 2, (16, 1)).float()
         with self.assertRaises(ValueError):
             loss_fn(input, target)
 
     def test_backward(self):
+        """Test that gradients can be computed."""
         loss_fn = AUCMLoss()
         input = torch.randn(32, 1, requires_grad=True)
         target = torch.randint(0, 2, (32, 1)).float()
@@ -70,6 +79,7 @@ def test_backward(self):
         self.assertIsNotNone(input.grad)
 
     def test_script_save(self):
+        """Test that the loss can be saved as TorchScript."""
         loss_fn = AUCMLoss()
         test_script_save(loss_fn, torch.randn(32, 1), torch.randint(0, 2, (32, 1)).float())