fix: handle constant B-spline mutual information

Signed-off-by: kyinhub <kevinpyin@gmail.com>

Author: yinkev

monai/losses/image_dissimilarity.py

@@ -214,7 +214,7 @@ def __init__(
                       IEEE Transactions in Medical Imaging. Vol.22, No.1,
                       January 2003. pp.120-128.
 
-            num_bins: number of bins for intensity
+            num_bins: number of bins for intensity. The b-spline kernel requires more than 4 bins.
             sigma_ratio: a hyper param for gaussian function
             reduction: {``"none"``, ``"mean"``, ``"sum"``}
                 Specifies the reduction to apply to the output. Defaults to ``"mean"``.
@@ -231,6 +231,8 @@ def __init__(
         bin_centers = torch.linspace(0.0, 1.0, num_bins)  # (num_bins,)
         sigma = torch.mean(bin_centers[1:] - bin_centers[:-1]) * sigma_ratio
         self.kernel_type = look_up_option(kernel_type, ["gaussian", "b-spline"])
+        if self.kernel_type == "b-spline" and num_bins <= 4:
+            raise ValueError(f"num_bins must be greater than 4 for b-spline kernel, got {num_bins}")
         self.num_bins = num_bins
         # declared as buffers so they move with the module (e.g. ``.to(device)``); only populated for the
         # gaussian kernel, hence the ``Tensor`` annotation reflects the type at the use sites in that path.
@@ -283,7 +285,13 @@ def parzen_windowing_b_spline(self, img: torch.Tensor, order: int) -> tuple[torc
         # window.
         _max, _min = torch.max(img), torch.min(img)
         padding = 2
-        bin_size = (_max - _min) / (self.num_bins - 2 * padding)
+        value_range = _max - _min
+        bin_size = value_range / (self.num_bins - 2 * padding)
+        bin_size = torch.where(
+            value_range > 0,
+            bin_size,
+            torch.ones_like(bin_size),
+        )
         norm_min = torch.div(_min, bin_size) - padding
 
         # assign bin/window index to each voxel
@@ -293,6 +301,8 @@ def parzen_windowing_b_spline(self, img: torch.Tensor, order: int) -> tuple[torc
         window_term = window_term.reshape(window_term.shape[0], -1, 1)  # (batch, num_sample, 1)
         bins = torch.arange(self.num_bins, device=window_term.device).reshape(1, 1, -1)  # (1, 1, num_bins)
         sample_bin_matrix = torch.abs(bins - window_term)  # (batch, num_sample, num_bins)
+        if sample_bin_matrix.dtype == torch.float16:
+            sample_bin_matrix = sample_bin_matrix.float()  # avoid overflow in the cubic polynomial
 
         # b-spleen kernel
         # (4 - 6 * abs ** 2 + 3 * abs ** 3) / 6 when 0 <= abs < 1

tests/losses/image_dissimilarity/test_global_mutual_information_loss.py

@@ -135,6 +135,10 @@ def test_b_spline_bin_centers_exists_as_none(self):
 
         self.assertIsNone(loss.bin_centers)
 
+    def test_b_spline_num_bins_must_allow_padding(self):
+        with self.assertRaisesRegex(ValueError, "num_bins must be greater than 4"):
+            GlobalMutualInformationLoss(kernel_type="b-spline", num_bins=4)
+
     @parameterized.expand(
         [
             (torch.ones((1, 2), dtype=torch.float), torch.ones((1, 3), dtype=torch.float)),  # mismatched_simple_dims
@@ -164,6 +168,34 @@ def test_ill_opts(self, num_bins, reduction, expected_exception, expected_messag
             GlobalMutualInformationLoss(num_bins=num_bins, reduction=reduction)(pred, target)
 
 
+class TestGlobalMutualInformationLossBSpline(unittest.TestCase):
+    def test_b_spline_constant_images_are_finite(self):
+        pred = torch.zeros((1, 1, 8, 8), requires_grad=True)
+        target = torch.ones_like(pred)
+        loss = GlobalMutualInformationLoss(kernel_type="b-spline")
+
+        result = loss(pred, target)
+
+        self.assertTrue(torch.isfinite(result))
+        self.assertAlmostEqual(result.item(), 0.0, places=6)
+        result.backward()
+        self.assertIsNotNone(pred.grad)
+        self.assertTrue(torch.isfinite(pred.grad).all())
+
+    def test_b_spline_constant_half_precision_images_are_finite(self):
+        pred = torch.zeros((1, 1, 8, 8), dtype=torch.float16, requires_grad=True)
+        target = torch.ones_like(pred)
+        loss = GlobalMutualInformationLoss(kernel_type="b-spline", num_bins=32)
+
+        result = loss(pred, target)
+
+        self.assertTrue(torch.isfinite(result))
+        self.assertAlmostEqual(result.item(), 0.0, places=6)
+        result.backward()
+        self.assertIsNotNone(pred.grad)
+        self.assertTrue(torch.isfinite(pred.grad).all())
+
+
 class TestGlobalMutualInformationLossBuffers(unittest.TestCase):
     def test_gaussian_kernel_registers_buffers(self):
         """Verify gaussian kernel registers preterm and bin_centers as non-trainable, non-persistent buffers."""