Repository: Project-MONAI/MONAI
File: tests/losses/test_focal_loss.py (Python)

Code before:
            FocalLoss(include_background=False, weight=(1.0, 1.0, -1.0))(chn_input, chn_target)

    def test_warnings(self):
        with self.assertWarns(Warning):
            chn_input = torch.ones((1, 1, 3))
            chn_target = torch.ones((1, 1, 3))
            loss = FocalLoss(to_onehot_y=True)
            loss(chn_input, chn_target)
        with self.assertWarns(Warning):
            chn_input = torch.ones((1, 1, 3))
            chn_target = torch.ones((1, 1, 3))
            loss = FocalLoss(include_background=False)
            loss(chn_input, chn_target)
        with self.assertWarns(Warning):
            chn_input = torch.ones((1, 3, 3))
            chn_target = torch.ones((1, 3, 3))
            loss = FocalLoss(include_background=False, use_softmax=True, alpha=0.5)
            loss(chn_input, chn_target)

    def test_script(self):
        for use_softmax in [True, False]:
            loss = FocalLoss(use_softmax=use_softmax)
            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

Review comment (suggestion):
Here I think we do need another test, or use `parameterized.expand`, to test with and without the background. 

Code after:
    def test_ill_class_weight(self):
        chn_input = torch.ones((1, 4, 3, 3))
        chn_target = torch.ones((1, 4, 3, 3))
        with self.assertRaisesRegex(ValueError, ""):
            FocalLoss(include_background=True, weight=(1.0, 1.0, 2.0))(chn_input, chn_target)
        with self.assertRaisesRegex(ValueError, ""):
            FocalLoss(include_background=False, weight=(1.0, 1.0, 1.0, 1.0))(chn_input, chn_target)
        with self.assertRaisesRegex(ValueError, ""):
            FocalLoss(include_background=False, weight=(1.0, 1.0, -1.0))(chn_input, chn_target)

    def test_warnings(self):
        with self.assertWarns(Warning):
            chn_input = torch.ones((1, 1, 3))
            chn_target = torch.ones((1, 1, 3))
            loss = FocalLoss(to_onehot_y=True)
            loss(chn_input, chn_target)
        with self.assertWarns(Warning):
            chn_input = torch.ones((1, 1, 3))
            chn_target = torch.ones((1, 1, 3))
            loss = FocalLoss(include_background=False)
            loss(chn_input, chn_target)
        with self.assertWarns(Warning):
            chn_input = torch.ones((1, 3, 3))
            chn_target = torch.ones((1, 3, 3))
            loss = FocalLoss(include_background=False, use_softmax=True, alpha=0.5)
            loss(chn_input, chn_target)

    def test_script(self):
        for use_softmax in [True, False]:
            loss = FocalLoss(use_softmax=use_softmax)
            test_input = torch.ones(2, 2, 8, 8)
            test_script_save(loss, test_input, test_input)

    @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
        batch_size = 2
        spatial_dims = (4, 4)

        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(