update test_unified_focal_loss.py

Signed-off-by: ytl0623 <[email protected]>

Author: ytl0623

monai/losses/unified_focal_loss.py

@@ -73,7 +73,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
 
         if not self.include_background:
             if n_pred_ch == 1:
-                warnings.warn("single channel prediction, `include_background=False` ignored.")
+                warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2)
             else:
                 # if skipping background, removing first channel
                 y_true = y_true[:, 1:]
@@ -110,8 +110,8 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
                 loss_list.append(1 - dice_class[:, i])
             else:
                 # Foreground classes: apply focal modulation
-                # Original logic: (1 - dice) * (1 - dice)^(-gamma) -> (1-dice)^(1-gamma)
-                loss_list.append((1 - dice_class[:, i]) * torch.pow(1 - dice_class[:, i], -self.gamma))
+                back_dice = torch.clamp(1 - dice_class[:, i], min=self.epsilon)
+                loss_list.append(back_dice * torch.pow(back_dice, -self.gamma))
 
         loss = torch.stack(loss_list, dim=-1)
 
@@ -150,11 +150,13 @@ def __init__(
             to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
             use_softmax: whether to use softmax to transform the original logits into probabilities.
                 If True, softmax is used. If False, sigmoid is used. Defaults to False.
+                softmax for mutually exclusive classes (standard multi-class).
+                sigmoid for multi-label/overlapping classes.
             reduction: Specifies the reduction to apply to the output. Defaults to ``"mean"``.
         """
         super().__init__(reduction=LossReduction(reduction).value)
         self.weight = weight
-        self.use_softmax = use_softmax  # 儲存參數
+        self.use_softmax = use_softmax
 
         self.focal_loss = FocalLoss(
             include_background=include_background,
@@ -175,8 +177,11 @@ def __init__(
     def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         """
         Args:
-            y_pred: (BNH[WD]) Logits (raw scores).
+            y_pred: (BNH[WD]) Logits (raw scores, not probabilities).
+                Do not pass pre-activated inputs; activation is applied internally.
             y_true: (BNH[WD]) Ground truth labels.
+        Returns:
+            torch.Tensor: Weighted combination of focal loss and asymmetric focal Tversky loss.
         """
         focal_loss = self.focal_loss(y_pred, y_true)
 

tests/losses/test_unified_focal_loss.py

@@ -19,18 +19,74 @@
 
 from monai.losses import AsymmetricUnifiedFocalLoss
 
+logit_pos = 10.0
+logit_neg = -10.0
+
 TEST_CASES = [
-    [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+    [  # Case 0: Binary segmentation
+        # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+        {
+            "use_softmax": False,
+            "include_background": True,
+        },
         {
-            "y_pred": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]),
-            "y_true": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]),
+            "y_pred": torch.tensor(
+                [[[[logit_pos, logit_neg], [logit_neg, logit_pos]]], [[[logit_pos, logit_neg], [logit_neg, logit_pos]]]]
+            ),
+            "y_true": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]], [[[1.0, 0.0], [0.0, 1.0]]]]),
         },
         0.0,
     ],
-    [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+    [  # Case 1: Multi-class segmentation with softmax
+        # shape: (1, 3, 2, 2), (1, 3, 2, 2)
         {
-            "y_pred": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]),
-            "y_true": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]),
+            "use_softmax": True,
+            "include_background": True,
+        },
+        {
+            "y_pred": torch.tensor(
+                [
+                    [
+                        [[logit_pos, logit_neg], [logit_neg, logit_neg]],  # Class 0 (background)
+                        [[logit_neg, logit_pos], [logit_neg, logit_neg]],  # Class 1
+                        [[logit_neg, logit_neg], [logit_pos, logit_pos]],  # Class 2
+                    ]
+                ]
+            ),
+            "y_true": torch.tensor(
+                [
+                    [
+                        [[1.0, 0.0], [0.0, 0.0]],  # Class 0 (background)
+                        [[0.0, 1.0], [0.0, 0.0]],  # Class 1
+                        [[0.0, 0.0], [1.0, 1.0]],  # Class 2
+                    ]
+                ]
+            ),
+        },
+        0.0,
+    ],
+    [  # Case 2: Multi-class with background excluded
+        # shape: (1, 3, 2, 2), (1, 3, 2, 2)
+        {"use_softmax": True, "include_background": False},
+        {
+            "y_pred": torch.tensor(
+                [
+                    [
+                        [[logit_pos, logit_neg], [logit_neg, logit_neg]],  # Class 0 (background)
+                        [[logit_neg, logit_pos], [logit_pos, logit_neg]],  # Class 1 (foreground)
+                        [[logit_neg, logit_neg], [logit_neg, logit_pos]],  # Class 2 (foreground)
+                    ]
+                ]
+            ),
+            "y_true": torch.tensor(
+                [
+                    [
+                        [[1.0, 0.0], [0.0, 0.0]],  # Class 0 (background)
+                        [[0.0, 1.0], [1.0, 0.0]],  # Class 1 (foreground)
+                        [[0.0, 0.0], [0.0, 1.0]],  # Class 2 (foreground)
+                    ]
+                ]
+            ),
         },
         0.0,
     ],
@@ -40,8 +96,16 @@
 class TestAsymmetricUnifiedFocalLoss(unittest.TestCase):
 
     @parameterized.expand(TEST_CASES)
-    def test_result(self, input_data, expected_val):
-        loss = AsymmetricUnifiedFocalLoss()
+    def test_result(self, input_param, input_data, expected_val):
+        """
+        Test AsymmetricUnifiedFocalLoss with various configurations.
+
+        Args:
+            input_param: Dict of loss constructor parameters (use_softmax, include_background, etc.).
+            input_data: Dict containing y_pred (logits) and y_true (ground truth) tensors.
+            expected_val: Expected loss value.
+        """
+        loss = AsymmetricUnifiedFocalLoss(**input_param)
         result = loss(**input_data)
         np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
 
@@ -52,8 +116,10 @@ def test_ill_shape(self):
 
     def test_with_cuda(self):
         loss = AsymmetricUnifiedFocalLoss()
-        i = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]])
-        j = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]])
+        i = torch.tensor(
+            [[[[logit_pos, logit_neg], [logit_neg, logit_pos]]], [[[logit_pos, logit_neg], [logit_neg, logit_pos]]]]
+        )
+        j = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]], [[[1.0, 0.0], [0.0, 1.0]]]])
         if torch.cuda.is_available():
             i = i.cuda()
             j = j.cuda()