diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 195ac32b1f..5457d636ae 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. @@ -281,13 +288,20 @@ def parzen_windowing_b_spline(self, img: torch.Tensor, order: int) -> tuple[torc # Note that there can still be non-zero bin values in the padded region, # it's just that these bins will never be a central bin for the Parzen # window. - _max, _min = torch.max(img), torch.min(img) + compute_img = img.float() if img.dtype in (torch.float16, torch.bfloat16) else img + _max, _min = torch.max(compute_img), torch.min(compute_img) padding = 2 - bin_size = (_max - _min) / (self.num_bins - 2 * padding) + interior_bins = self.num_bins - 2 * padding + value_range = _max - _min + raw_bin_size = value_range / interior_bins + # Smaller bins either underflow the compute dtype or require a + # normalization slope that cannot be represented by the input dtype. + minimum_bin_size = max(torch.finfo(compute_img.dtype).tiny, 1.0 / torch.finfo(img.dtype).max) + bin_size = torch.where(raw_bin_size >= minimum_bin_size, raw_bin_size, torch.ones_like(raw_bin_size)) norm_min = torch.div(_min, bin_size) - padding # assign bin/window index to each voxel - window_term = torch.div(img, bin_size) - norm_min # B[NDHW] + window_term = torch.div(compute_img, bin_size) - norm_min # B[NDHW] # make sure the extreme values are in valid (non-padded) bins window_term = torch.clamp(window_term, padding, self.num_bins - padding - 1) # B[NDHW] window_term = window_term.reshape(window_term.shape[0], -1, 1) # (batch, num_sample, 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..999d4d707f 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,65 @@ 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.""" + + @parameterized.expand(["prediction", "target"]) + def test_b_spline_single_constant_input_is_finite(self, constant_input): + """Verify either independently constant input yields finite gradients.""" + varying = torch.linspace(0.0, 1.0, 64).reshape(1, 1, 8, 8) + if constant_input == "prediction": + pred = torch.zeros_like(varying, requires_grad=True) + target = varying + else: + pred = varying.clone().requires_grad_() + target = torch.ones_like(varying) + loss = GlobalMutualInformationLoss(kernel_type="b-spline") + + result = loss(pred, target) + + self.assertTrue(torch.isfinite(result)) + 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 survive overflow-sized bin distances.""" + 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=64) + + 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()) + + @parameterized.expand( + [ + ("float16_tiny", torch.float16, torch.finfo(torch.float16).tiny / 2), + ("bfloat16_tiny", torch.bfloat16, torch.finfo(torch.bfloat16).tiny / 2), + ("float32_tiny", torch.float32, torch.finfo(torch.float32).tiny / 2), + ("float16_large", torch.float16, 65000.0), + ] + ) + def test_b_spline_nonzero_ranges_are_finite(self, _, dtype, maximum): + """Verify extreme nonzero ranges yield finite loss and gradients.""" + values = torch.tensor([0.0, maximum, maximum, 0.0], dtype=dtype).reshape(1, 1, 2, 2) + pred = values.clone().requires_grad_() + target = torch.flip(values, dims=(-1,)) + loss = GlobalMutualInformationLoss(kernel_type="b-spline", num_bins=64) + + result = loss(pred, target) + + self.assertTrue(torch.isfinite(result)) + 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."""