update test_unified_focal_loss.py

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

Author: ytl0623

monai/losses/unified_focal_loss.py

@@ -111,7 +111,8 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             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)
 
@@ -154,7 +155,7 @@ def __init__(
         """
         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,

tests/losses/test_unified_focal_loss.py

@@ -19,18 +19,42 @@
 
 from monai.losses import AsymmetricUnifiedFocalLoss
 
+# Helper to create high confidence logits (approx 10 -> sigmoid close to 1, -10 -> sigmoid close to 0)
+logit_pos = 10.0
+logit_neg = -10.0
+
 TEST_CASES = [
-    [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+    # Case 0: Binary Segmentation (Sigmoid), Perfect Prediction
+    # Shape: (B=2, C=1, H=2, W=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]]]]),
+            # Logits: High value where ground truth is 1, Low value where ground truth is 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 (Softmax), Perfect Prediction
+    # Shape: (B=1, C=3, H=1, W=2) -> 3 classes (0: Background, 1: Class A, 2: Class B)
+    [
+        {
+            "use_softmax": True,
+            "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]]]]),
+            # Logits:
+            # Pixel 1: Class 0 is target -> [10, -10, -10]
+            # Pixel 2: Class 2 is target -> [-10, -10, 10]
+            "y_pred": torch.tensor(
+                [[[[logit_pos, logit_neg], [logit_neg, logit_neg], [logit_neg, logit_pos]]]]  # Ch 0  # Ch 1  # Ch 2
+            ),
+            "y_true": torch.tensor([[[[1.0, 0.0], [0.0, 0.0], [0.0, 1.0]]]]),  # Ch 0 (Background)  # Ch 1  # Ch 2
         },
         0.0,
     ],
@@ -40,8 +64,8 @@
 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):
+        loss = AsymmetricUnifiedFocalLoss(**input_param)
         result = loss(**input_data)
         np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
 
@@ -51,12 +75,15 @@ def test_ill_shape(self):
             loss(torch.ones((2, 2, 2)), torch.ones((2, 2, 2, 2)))
 
     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]]]])
+        loss = AsymmetricUnifiedFocalLoss(use_softmax=False)
+        i = torch.tensor([[[[logit_pos, logit_neg], [logit_neg, logit_pos]]]])
+        j = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]])
+
         if torch.cuda.is_available():
             i = i.cuda()
             j = j.cuda()
+            loss = loss.cuda()
+
         output = loss(i, j)
         print(output)
         np.testing.assert_allclose(output.detach().cpu().numpy(), 0.0, atol=1e-4, rtol=1e-4)