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

Code before:
            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
            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")


Review comment (suggestion):
We have [this definition](https://github.com/Project-MONAI/MONAI/blob/01711cf7e3af991aabc92cc3fdf4c0b65949e65e/tests/test_utils.py#L874C1-L874C14) to use in place of this line if you want. Another approach is to have a `device` argument and use the `parameterized.expand` decorator [like here](https://github.com/Project-MONAI/MONAI/blob/01711cf7e3af991aabc92cc3fdf4c0b65949e65e/tests/data/meta_tensor/test_meta_tensor.py#L201).

Code after:

    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(
            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)