add test case for alpha as a sequence

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

Author: ytl0623

monai/losses/focal_loss.py

@@ -165,18 +165,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         input = input.float()
         target = target.float()
 
-        alpha_arg: float | torch.Tensor | None = self.alpha
-        if isinstance(alpha_arg, torch.Tensor):
-            alpha_arg = alpha_arg.to(input.device)
-
         if self.use_softmax:
             if not self.include_background and self.alpha is not None:
                 if isinstance(self.alpha, (float, int)):
-                    alpha_arg = None
                     warnings.warn("`include_background=False`, scalar `alpha` ignored when using softmax.")
-            loss = softmax_focal_loss(input, target, self.gamma, alpha_arg)
+            loss = softmax_focal_loss(input, target, self.gamma, self.alpha)
         else:
-            loss = sigmoid_focal_loss(input, target, self.gamma, alpha_arg)
+            loss = sigmoid_focal_loss(input, target, self.gamma, self.alpha)
 
         num_of_classes = target.shape[1]
         if self.class_weight is not None and num_of_classes != 1:

tests/losses/test_focal_loss.py

@@ -374,6 +374,47 @@ 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):
+        """
+        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)
+
 
 if __name__ == "__main__":
     unittest.main()