diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 195ac32b1f..be36460647 100644 --- a/monai/losses/image_dissimilarity.py +++ b/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"``. @@ -224,6 +224,10 @@ def __init__( - ``"sum"``: the output will be summed. smooth_nr: a small constant added to the numerator to avoid nan. smooth_dr: a small constant added to the denominator to avoid nan. + + Raises: + ValueError: if ``num_bins`` is not positive, or if a B-spline + kernel is configured with four or fewer bins. """ super().__init__(reduction=LossReduction(reduction).value) if num_bins <= 0: @@ -231,6 +235,9 @@ 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"]) + # B-spline windowing reserves two padding bins at each boundary. + 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 +290,9 @@ 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 +302,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 diff --git a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py index 19a60f7219..4735da7378 100644 --- a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py +++ b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py @@ -135,6 +135,11 @@ 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): + """Verify B-spline bin counts leave room for boundary padding.""" + 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 +169,38 @@ 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): + """Test B-spline mutual information on degenerate intensity ranges.""" + + def test_b_spline_constant_images_are_finite(self): + """Verify constant float32 inputs yield finite loss and gradients.""" + 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): + """Verify constant float16 inputs yield finite loss and gradients.""" + 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."""