another test to test without the background

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

Author: ytl0623

monai/losses/focal_loss.py

@@ -168,7 +168,9 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         if self.use_softmax:
             if not self.include_background and self.alpha is not None:
                 if isinstance(self.alpha, (float, int)):
-                    warnings.warn("`include_background=False`, scalar `alpha` ignored when using softmax.")
+                    warnings.warn(
+                        "`include_background=False`, scalar `alpha` ignored when using softmax.", stacklevel=2
+                    )
             loss = softmax_focal_loss(input, target, self.gamma, self.alpha)
         else:
             loss = sigmoid_focal_loss(input, target, self.gamma, self.alpha)

tests/losses/test_focal_loss.py

@@ -21,7 +21,7 @@
 
 from monai.losses import FocalLoss
 from monai.networks import one_hot
-from tests.test_utils import test_script_save
+from tests.test_utils import test_script_save, TEST_DEVICES
 
 TEST_CASES = []
 for device in ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]:
@@ -77,6 +77,13 @@
     TEST_CASES.append([{"to_onehot_y": True, "use_softmax": True}, input_data, 0.16276])
     TEST_CASES.append([{"to_onehot_y": True, "alpha": 0.8, "use_softmax": True}, input_data, 0.08138])
 
+TEST_ALPHA_BROADCASTING = []
+for case in TEST_DEVICES:
+    device = case[0]
+    for include_background in [True, False]:
+        for use_softmax in [True, False]:
+            TEST_ALPHA_BROADCASTING.append([device, include_background, use_softmax])
+
 
 class TestFocalLoss(unittest.TestCase):
     @parameterized.expand(TEST_CASES)
@@ -374,46 +381,39 @@ def test_script(self):
             test_input = torch.ones(2, 2, 8, 8)
             test_script_save(loss, test_input, test_input)
 
-    def test_alpha_sequence_broadcasting(self):
+    @parameterized.expand(TEST_ALPHA_BROADCASTING)
+    def test_alpha_sequence_broadcasting(self, device, include_background, use_softmax):
         """
         Test FocalLoss with alpha as a sequence for proper broadcasting.
         """
         num_classes = 3
-        alpha_seq = [0.1, 0.5, 2.0]
         batch_size = 2
         spatial_dims = (4, 4)
 
-        devices = ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]
-
-        for device in devices:
-            logits = torch.randn(batch_size, num_classes, *spatial_dims, device=device)
-            target = torch.randint(0, num_classes, (batch_size, 1, *spatial_dims), device=device)
-
-            # Case 1: Softmax + Alpha Sequence
-            loss_func_softmax = FocalLoss(
-                to_onehot_y=True, gamma=2.0, alpha=alpha_seq, use_softmax=True, reduction="mean"
-            )
-            loss_soft = loss_func_softmax(logits, target)
-
-            self.assertTrue(torch.is_tensor(loss_soft))
-            self.assertEqual(loss_soft.ndim, 0)
-            self.assertTrue(loss_soft > 0, f"Softmax loss on {device} should be positive")
-
-            # Case 2: Sigmoid + Alpha Sequence
-            loss_func_sigmoid = FocalLoss(
-                to_onehot_y=True, gamma=2.0, alpha=alpha_seq, use_softmax=False, reduction="mean"
-            )
-            loss_sig = loss_func_sigmoid(logits, target)
-
-            self.assertTrue(torch.is_tensor(loss_sig))
-            self.assertEqual(loss_sig.ndim, 0)
-            self.assertTrue(loss_sig > 0, f"Sigmoid loss on {device} should be positive")
-
-            # Case 3: Error Handling (Mismatched alpha length)
-            if device == devices[0]:
-                wrong_alpha = [0.1, 0.5]
-                with self.assertRaisesRegex(ValueError, "length of alpha"):
-                    FocalLoss(to_onehot_y=True, alpha=wrong_alpha, use_softmax=True)(logits, target)
+        logits = torch.randn(batch_size, num_classes, *spatial_dims, device=device)
+        target = torch.randint(0, num_classes, (batch_size, 1, *spatial_dims), device=device)
+
+        if include_background:
+            alpha_seq = [0.1, 0.5, 2.0]
+        else:
+            alpha_seq = [0.5, 2.0]
+
+        loss_func = FocalLoss(
+            to_onehot_y=True,
+            gamma=2.0,
+            alpha=alpha_seq,
+            include_background=include_background,
+            use_softmax=use_softmax,
+            reduction="mean",
+        )
+
+        result = loss_func(logits, target)
+
+        self.assertTrue(torch.is_tensor(result))
+        self.assertEqual(result.ndim, 0)
+        self.assertTrue(
+            result > 0, f"Loss should be positive. params: dev={device}, bg={include_background}, softmax={use_softmax}"
+        )
 
 
 if __name__ == "__main__":