update test_unified_focal_loss.py

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

Author: ytl0623

tests/losses/test_unified_focal_loss.py

@@ -12,38 +12,62 @@
 from __future__ import annotations
 
 import unittest
-
 import numpy as np
 import torch
 from parameterized import parameterized
-
 from monai.losses import AsymmetricUnifiedFocalLoss
 
+LOGIT_HIGH = 5.0
+LOGIT_LOW = -5.0
+
 TEST_CASES = [
-    [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+    # Case 0: Sigmoid + Include BG
+    [
+        {"use_softmax": False, "include_background": True},
+        {
+            "y_pred": torch.tensor([[[[LOGIT_HIGH]], [[LOGIT_LOW]]]]),
+            "y_true": torch.tensor([[[[1.0]], [[0.0]]]]),
+        },
+    ],
+    # Case 1: Softmax + Ignore BG
+    [
+        {"use_softmax": True, "include_background": False},
+        {
+            "y_pred": torch.tensor([[[[LOGIT_LOW]], [[LOGIT_HIGH]], [[LOGIT_LOW]]]]), 
+            "y_true": torch.tensor([[[[0.0]], [[1.0]], [[0.0]]]]),
+        },
+    ],
+    # Case 2: Softmax + Include BG
+    [
+        {"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]]]]),
+            "y_pred": torch.tensor([[[[LOGIT_HIGH]], [[LOGIT_LOW]], [[LOGIT_LOW]]]]), 
+            "y_true": torch.tensor([[[[1.0]], [[0.0]], [[0.0]]]]),
         },
-        0.0,
     ],
-    [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+    # Case 3: Sigmoid + Ignore BG
+    [
+        {"use_softmax": False, "include_background": False},
         {
-            "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]]]]),
+            "y_pred": torch.tensor([[[[LOGIT_HIGH]], [[LOGIT_HIGH]]]]),
+            "y_true": torch.tensor([[[[0.0]], [[1.0]]]]),
         },
-        0.0,
     ],
 ]
 
-
 class TestAsymmetricUnifiedFocalLoss(unittest.TestCase):
 
     @parameterized.expand(TEST_CASES)
-    def test_result(self, input_data, expected_val):
-        loss = AsymmetricUnifiedFocalLoss()
-        result = loss(**input_data)
-        np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
+    def test_result(self, input_param, input_data):
+        loss_func = AsymmetricUnifiedFocalLoss(**input_param)
+        result = loss_func(**input_data)
+        res_val = result.detach().cpu().numpy()
+        
+        print(f"Params: {input_param} -> Loss: {res_val}")
+
+
+        self.assertFalse(np.isnan(res_val), "Loss should not be NaN")
+        self.assertTrue(res_val < 1.0, f"Loss {res_val} is too high (expected < 1.0)")
 
     def test_ill_shape(self):
         loss = AsymmetricUnifiedFocalLoss()
@@ -52,15 +76,19 @@ def test_ill_shape(self):
 
     def test_with_cuda(self):
         loss = AsymmetricUnifiedFocalLoss()
-        i = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]])
+        i = torch.tensor([[[[5.0, -5.0], [-5.0, 5.0]]], [[[5.0, -5.0], [-5.0, 5.0]]]])
         j = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.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)
-
+        res_val = output.detach().cpu().numpy()
+        
+        self.assertFalse(np.isnan(res_val), "CUDA Loss should not be NaN")
+        self.assertTrue(res_val < 1.0, f"CUDA Loss {res_val} is too high")
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()