From 27eb62d3931aedcbab491feb79e1edeac8afabdb Mon Sep 17 00:00:00 2001 From: kyinhub Date: Sat, 25 Jul 2026 11:18:31 -0700 Subject: [PATCH 1/5] fix(losses): avoid half precision mutual information overflow Signed-off-by: kyinhub --- monai/losses/image_dissimilarity.py | 59 ++++++++++++++----- .../test_global_mutual_information_loss.py | 53 +++++++++++++++++ 2 files changed, 97 insertions(+), 15 deletions(-) diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 195ac32b1f..26fa017077 100644 --- a/monai/losses/image_dissimilarity.py +++ b/monai/losses/image_dissimilarity.py @@ -316,16 +316,30 @@ def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, to Note: the input is expected to range between 0 and 1 Args: img: the shape should be B[NDHW]. + + Returns: + A tuple containing per-sample Gaussian bin weights and the + corresponding marginal probability. + + Raises: + ValueError: if the Gaussian kernel buffers are unavailable. """ - img = torch.clamp(img, 0, 1) + output_dtype = img.dtype + if output_dtype == torch.float16: + img = img.float() + compute_dtype = torch.float32 if img.dtype in (torch.float16, torch.bfloat16) else img.dtype + img = torch.clamp(img, 0, 1).to(dtype=compute_dtype) img = img.reshape(img.shape[0], -1, 1) # (batch, num_sample, 1) if self.bin_centers is None or self.preterm is None: raise ValueError("bin_centers and preterm must be defined for gaussian parzen windowing.") - weight = torch.exp( - -self.preterm.to(img) * (img - self.bin_centers.to(img)) ** 2 - ) # (batch, num_sample, num_bin) + preterm = self.preterm.to(device=img.device, dtype=compute_dtype) + bin_centers = self.bin_centers.to(device=img.device, dtype=compute_dtype) + weight = torch.exp(-preterm * (img - bin_centers) ** 2) # (batch, num_sample, num_bin) weight = weight / torch.sum(weight, dim=-1, keepdim=True) # (batch, num_sample, num_bin) probability = torch.mean(weight, dim=-2, keepdim=True) # (batch, 1, num_bin) + if output_dtype in (torch.float16, torch.bfloat16): + weight = weight.to(dtype=output_dtype) + probability = probability.to(dtype=output_dtype) return weight, probability def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: @@ -333,6 +347,9 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: Args: pred: the shape should be B[NDHW]. target: the shape should be same as the pred shape. + Returns: + Reduced negative mutual information loss. + Raises: ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. """ @@ -340,16 +357,28 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: raise ValueError(f"ground truth has differing shape ({target.shape}) from pred ({pred.shape})") wa, pa, wb, pb = self.parzen_windowing(pred, target) # (batch, num_sample, num_bin), (batch, 1, num_bin) - pab = torch.bmm(wa.permute(0, 2, 1), wb.to(wa)).div(wa.shape[1]) # (batch, num_bins, num_bins) - papb = torch.bmm(pa.permute(0, 2, 1), pb.to(pa)) # (batch, num_bins, num_bins) - mi = torch.sum( - pab * torch.log((pab + self.smooth_nr) / (papb + self.smooth_dr) + self.smooth_dr), dim=(1, 2) - ) # (batch) + # Half-precision matrix multiplication can overflow before the joint + # histogram is divided by the number of samples. Accumulate histogram + # products in float32 while preserving float64 inputs. + output_dtype = wa.dtype + compute_dtype = torch.float32 if wa.dtype in (torch.float16, torch.bfloat16) else wa.dtype + wa = wa.to(dtype=compute_dtype) + wb = wb.to(wa) + pa = pa.to(dtype=compute_dtype) + pb = pb.to(pa) + with torch.autocast(device_type=wa.device.type, enabled=False): + pab = torch.bmm(wa.permute(0, 2, 1), wb).div(wa.shape[1]) # (batch, num_bins, num_bins) + papb = torch.bmm(pa.permute(0, 2, 1), pb) # (batch, num_bins, num_bins) + mi = torch.sum( + pab * torch.log((pab + self.smooth_nr) / (papb + self.smooth_dr) + self.smooth_dr), dim=(1, 2) + ) # (batch) if self.reduction == LossReduction.SUM.value: - return torch.sum(mi).neg() # sum over the batch and channel ndims - if self.reduction == LossReduction.NONE.value: - return mi.neg() - if self.reduction == LossReduction.MEAN.value: - return torch.mean(mi).neg() # average over the batch and channel ndims - raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].') + loss = torch.sum(mi).neg() # sum over the batch and channel ndims + elif self.reduction == LossReduction.NONE.value: + loss = mi.neg() + elif self.reduction == LossReduction.MEAN.value: + loss = torch.mean(mi).neg() # average over the batch and channel ndims + else: + raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].') + return loss.to(dtype=output_dtype) 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..bb0ba65973 100644 --- a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py +++ b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py @@ -164,6 +164,59 @@ def test_ill_opts(self, num_bins, reduction, expected_exception, expected_messag GlobalMutualInformationLoss(num_bins=num_bins, reduction=reduction)(pred, target) +class TestGlobalMutualInformationLossHalfPrecision(unittest.TestCase): + """Test stable Gaussian mutual information in reduced-precision modes.""" + + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_half_precision_gaussian_weights_with_many_bins_are_finite(self, dtype): + """Verify many-bin Parzen outputs remain finite and preserve metadata.""" + image = torch.zeros((1, 1, 2), dtype=dtype) + loss = GlobalMutualInformationLoss(kernel_type="gaussian", num_bins=256) + + weight, probability = loss.parzen_windowing_gaussian(image) + + self.assertTrue(torch.isfinite(weight).all()) + self.assertTrue(torch.isfinite(probability).all()) + self.assertEqual(weight.dtype, image.dtype) + self.assertEqual(probability.dtype, image.dtype) + self.assertEqual(weight.device, image.device) + self.assertEqual(probability.device, image.device) + + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_half_precision_large_constant_volume_is_finite(self, dtype): + """Verify reduced-precision loss and gradients remain finite.""" + pred = torch.zeros((1, 1, 48, 48, 48), dtype=dtype, requires_grad=True) + target = torch.zeros_like(pred) + loss = GlobalMutualInformationLoss(kernel_type="gaussian") + + result = loss(pred, target) + + self.assertTrue(torch.isfinite(result)) + self.assertEqual(result.dtype, pred.dtype) + self.assertEqual(result.device, pred.device) + result.backward() + self.assertIsNotNone(pred.grad) + self.assertTrue(torch.isfinite(pred.grad).all()) + self.assertEqual(pred.grad.dtype, pred.dtype) + self.assertEqual(pred.grad.device, pred.device) + + def test_cpu_float16_autocast_large_volume_is_finite(self): + """Verify CPU float16 autocast avoids histogram accumulation overflow.""" + pred = torch.zeros((1, 1, 48, 48, 48), requires_grad=True) + target = torch.zeros_like(pred) + loss = GlobalMutualInformationLoss(kernel_type="gaussian") + + with torch.autocast(device_type="cpu", dtype=torch.float16): + result = loss(pred, target) + + self.assertTrue(torch.isfinite(result)) + result.backward() + self.assertIsNotNone(pred.grad) + self.assertTrue(torch.isfinite(pred.grad).all()) + self.assertEqual(pred.grad.dtype, pred.dtype) + self.assertEqual(pred.grad.device, pred.device) + + 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.""" From 76a50a8443d8c8c5b7700d541b6c0cb7cff2328e Mon Sep 17 00:00:00 2001 From: kyinhub Date: Sat, 25 Jul 2026 14:04:37 -0700 Subject: [PATCH 2/5] fix(losses): strengthen reduced-precision MI contracts Signed-off-by: kyinhub --- monai/losses/image_dissimilarity.py | 56 ++++++--- .../test_global_mutual_information_loss.py | 106 ++++++++++++++++++ 2 files changed, 146 insertions(+), 16 deletions(-) diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 26fa017077..0d80801f1a 100644 --- a/monai/losses/image_dissimilarity.py +++ b/monai/losses/image_dissimilarity.py @@ -11,6 +11,8 @@ from __future__ import annotations +import math + import torch from torch.nn import functional as F from torch.nn.modules.loss import _Loss @@ -236,10 +238,18 @@ def __init__( # gaussian kernel, hence the ``Tensor`` annotation reflects the type at the use sites in that path. self.preterm: torch.Tensor | None self.bin_centers: torch.Tensor | None + self._preterm_value: float | None = None self.register_buffer("preterm", None, persistent=False) self.register_buffer("bin_centers", None, persistent=False) if self.kernel_type == "gaussian": - self.register_buffer("preterm", 1 / (2 * sigma**2), persistent=False) + preterm = 1 / (2 * sigma**2) + preterm_value = float(preterm) + if not math.isfinite(preterm_value) and num_bins > 1 and sigma_ratio != 0.0: + preterm_value = (num_bins - 1) ** 2 / (2.0 * sigma_ratio**2) + if not bool(torch.isfinite(preterm)): + preterm = torch.as_tensor(preterm_value, dtype=torch.float32) + self._preterm_value = preterm_value + self.register_buffer("preterm", preterm, persistent=False) self.register_buffer("bin_centers", bin_centers[None, None, ...], persistent=False) self.smooth_nr = float(smooth_nr) self.smooth_dr = float(smooth_dr) @@ -325,14 +335,17 @@ def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, to ValueError: if the Gaussian kernel buffers are unavailable. """ output_dtype = img.dtype - if output_dtype == torch.float16: - img = img.float() - compute_dtype = torch.float32 if img.dtype in (torch.float16, torch.bfloat16) else img.dtype - img = torch.clamp(img, 0, 1).to(dtype=compute_dtype) + compute_dtype = torch.float32 if output_dtype in (torch.float16, torch.bfloat16) else output_dtype + img = torch.clamp(img.to(dtype=compute_dtype), 0, 1) img = img.reshape(img.shape[0], -1, 1) # (batch, num_sample, 1) - if self.bin_centers is None or self.preterm is None: + if self.bin_centers is None or self.preterm is None or self._preterm_value is None: raise ValueError("bin_centers and preterm must be defined for gaussian parzen windowing.") preterm = self.preterm.to(device=img.device, dtype=compute_dtype) + preterm = torch.where( + torch.isfinite(preterm), + preterm, + torch.as_tensor(self._preterm_value, device=img.device, dtype=compute_dtype), + ) bin_centers = self.bin_centers.to(device=img.device, dtype=compute_dtype) weight = torch.exp(-preterm * (img - bin_centers) ** 2) # (batch, num_sample, num_bin) weight = weight / torch.sum(weight, dim=-1, keepdim=True) # (batch, num_sample, num_bin) @@ -351,27 +364,38 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: Reduced negative mutual information loss. Raises: - ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. + ValueError: if ``pred`` and ``target`` have different shapes, or + if ``self.reduction`` is not one of ``"mean"``, ``"sum"``, + or ``"none"``. """ if target.shape != pred.shape: raise ValueError(f"ground truth has differing shape ({target.shape}) from pred ({pred.shape})") wa, pa, wb, pb = self.parzen_windowing(pred, target) # (batch, num_sample, num_bin), (batch, 1, num_bin) - # Half-precision matrix multiplication can overflow before the joint - # histogram is divided by the number of samples. Accumulate histogram - # products in float32 while preserving float64 inputs. + # A half-precision matrix product can overflow while accumulating the + # unnormalized joint histogram. Eager execution disables autocast for + # this operation. TorchScript cannot compile a dynamic autocast device, + # so it computes the normalized histogram directly by scaling both + # operands by sqrt(N). output_dtype = wa.dtype compute_dtype = torch.float32 if wa.dtype in (torch.float16, torch.bfloat16) else wa.dtype wa = wa.to(dtype=compute_dtype) wb = wb.to(wa) pa = pa.to(dtype=compute_dtype) pb = pb.to(pa) - with torch.autocast(device_type=wa.device.type, enabled=False): - pab = torch.bmm(wa.permute(0, 2, 1), wb).div(wa.shape[1]) # (batch, num_bins, num_bins) - papb = torch.bmm(pa.permute(0, 2, 1), pb) # (batch, num_bins, num_bins) - mi = torch.sum( - pab * torch.log((pab + self.smooth_nr) / (papb + self.smooth_dr) + self.smooth_dr), dim=(1, 2) - ) # (batch) + if torch.jit.is_scripting(): + sample_scale = float(wa.shape[1]) ** 0.5 + pab = torch.bmm((wa / sample_scale).permute(0, 2, 1), wb / sample_scale).to( + dtype=compute_dtype + ) # (batch, num_bins, num_bins) + papb = torch.bmm(pa.permute(0, 2, 1), pb).to(dtype=compute_dtype) # (batch, num_bins, num_bins) + else: + with torch.autocast(device_type=wa.device.type, enabled=False): + pab = torch.bmm(wa.permute(0, 2, 1), wb).div(wa.shape[1]) # (batch, num_bins, num_bins) + papb = torch.bmm(pa.permute(0, 2, 1), pb) # (batch, num_bins, num_bins) + mi = torch.sum( + pab * torch.log((pab + self.smooth_nr) / (papb + self.smooth_dr) + self.smooth_dr), dim=(1, 2) + ) # (batch) if self.reduction == LossReduction.SUM.value: loss = torch.sum(mi).neg() # sum over the batch and channel ndims 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 bb0ba65973..d7ebd9b90c 100644 --- a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py +++ b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py @@ -181,6 +181,80 @@ def test_half_precision_gaussian_weights_with_many_bins_are_finite(self, dtype): self.assertEqual(probability.dtype, image.dtype) self.assertEqual(weight.device, image.device) self.assertEqual(probability.device, image.device) + torch.testing.assert_close( + weight.float().sum(dim=-1), torch.ones_like(weight[..., 0], dtype=torch.float32), rtol=0.0, atol=5e-3 + ) + torch.testing.assert_close( + probability.float().sum(dim=-1), + torch.ones_like(probability[..., 0], dtype=torch.float32), + rtol=0.0, + atol=5e-3, + ) + + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_module_cast_with_many_bins_remains_finite(self, dtype): + """Verify module dtype conversion cannot overflow Gaussian parameters. + + Args: + dtype: reduced-precision floating-point dtype to test. + """ + image = torch.linspace(0.0, 1.0, 64, dtype=dtype).reshape(1, 1, 8, 8).requires_grad_() + target = torch.flip(image.detach(), dims=(-1,)) + loss = GlobalMutualInformationLoss(kernel_type="gaussian", num_bins=256).to(dtype=dtype) + + weight, probability = loss.parzen_windowing_gaussian(image) + result = loss(image, target) + + self.assertTrue(torch.isfinite(weight).all()) + self.assertTrue(torch.isfinite(probability).all()) + self.assertTrue(torch.isfinite(result)) + result.backward() + self.assertIsNotNone(image.grad) + self.assertTrue(torch.isfinite(image.grad).all()) + + def test_float16_default_dtype_with_many_bins_remains_finite(self): + """Verify construction under a float16 default keeps Gaussian parameters finite.""" + original_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.float16) + image = torch.linspace(0.0, 1.0, 64).reshape(1, 1, 8, 8).requires_grad_() + target = torch.flip(image.detach(), dims=(-1,)) + loss = GlobalMutualInformationLoss(kernel_type="gaussian", num_bins=256) + + weight, probability = loss.parzen_windowing_gaussian(image) + result = loss(image, target) + + self.assertTrue(torch.isfinite(weight).all()) + self.assertTrue(torch.isfinite(probability).all()) + self.assertTrue(torch.isfinite(result)) + result.backward() + self.assertIsNotNone(image.grad) + self.assertTrue(torch.isfinite(image.grad).all()) + finally: + torch.set_default_dtype(original_dtype) + + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_half_precision_nonconstant_images_match_float32(self, dtype): + """Verify nonconstant reduced-precision loss tracks float32. + + Args: + dtype: reduced-precision floating-point dtype to test. + """ + pred_float = torch.linspace(0.0, 1.0, 64).reshape(1, 1, 8, 8) + target_float = torch.flip(pred_float, dims=(-1,)) + loss = GlobalMutualInformationLoss(kernel_type="gaussian") + expected = loss(pred_float, target_float) + pred = pred_float.to(dtype=dtype).requires_grad_() + target = target_float.to(dtype=dtype) + + result = loss(pred, target) + + self.assertTrue(torch.isfinite(result)) + self.assertEqual(result.dtype, dtype) + torch.testing.assert_close(result.float(), expected, rtol=1e-2, atol=1e-2) + result.backward() + self.assertIsNotNone(pred.grad) + self.assertTrue(torch.isfinite(pred.grad).all()) @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) def test_half_precision_large_constant_volume_is_finite(self, dtype): @@ -200,6 +274,37 @@ def test_half_precision_large_constant_volume_is_finite(self, dtype): self.assertEqual(pred.grad.dtype, pred.dtype) self.assertEqual(pred.grad.device, pred.device) + def test_cpu_float16_autocast_nonconstant_images_match_float32(self): + """Verify nonconstant CPU autocast loss matches float32.""" + pred = torch.linspace(0.0, 1.0, 64).reshape(1, 1, 8, 8).requires_grad_() + target = torch.flip(pred.detach(), dims=(-1,)) + loss = GlobalMutualInformationLoss(kernel_type="gaussian") + expected = loss(pred, target).detach() + + with torch.autocast(device_type="cpu", dtype=torch.float16): + result = loss(pred, target) + + self.assertTrue(torch.isfinite(result)) + self.assertEqual(result.dtype, pred.dtype) + torch.testing.assert_close(result, expected) + result.backward() + self.assertIsNotNone(pred.grad) + self.assertTrue(torch.isfinite(pred.grad).all()) + + def test_scripted_cpu_float16_autocast_large_volume_is_finite(self): + """Verify scripted loss avoids float16 histogram overflow under autocast.""" + pred = torch.zeros((1, 1, 257, 257), requires_grad=True) + target = torch.zeros_like(pred) + loss = torch.jit.script(GlobalMutualInformationLoss(kernel_type="gaussian")) + + with torch.autocast(device_type="cpu", dtype=torch.float16): + result = loss(pred, target) + + self.assertTrue(torch.isfinite(result)) + result.backward() + self.assertIsNotNone(pred.grad) + self.assertTrue(torch.isfinite(pred.grad).all()) + def test_cpu_float16_autocast_large_volume_is_finite(self): """Verify CPU float16 autocast avoids histogram accumulation overflow.""" pred = torch.zeros((1, 1, 48, 48, 48), requires_grad=True) @@ -210,6 +315,7 @@ def test_cpu_float16_autocast_large_volume_is_finite(self): result = loss(pred, target) self.assertTrue(torch.isfinite(result)) + self.assertEqual(result.dtype, pred.dtype) result.backward() self.assertIsNotNone(pred.grad) self.assertTrue(torch.isfinite(pred.grad).all()) From f884e59c6a4fe02dfcf6c2991c6dfcf7ba32dc9b Mon Sep 17 00:00:00 2001 From: kyinhub Date: Wed, 29 Jul 2026 22:22:20 -0700 Subject: [PATCH 3/5] fix(losses): retain Gaussian Parzen compute precision Signed-off-by: kyinhub --- monai/losses/image_dissimilarity.py | 35 +++++++++++++++---- .../test_global_mutual_information_loss.py | 23 ++++++++++++ 2 files changed, 51 insertions(+), 7 deletions(-) diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 0d80801f1a..0f4d658b61 100644 --- a/monai/losses/image_dissimilarity.py +++ b/monai/losses/image_dissimilarity.py @@ -255,11 +255,26 @@ def __init__( self.smooth_dr = float(smooth_dr) def parzen_windowing( - self, pred: torch.Tensor, target: torch.Tensor + self, pred: torch.Tensor, target: torch.Tensor, restore_input_dtype: bool = True ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Apply the configured Parzen window to both inputs. + + Args: + pred: the prediction tensor. + target: the target tensor. + restore_input_dtype: whether Gaussian weights and probabilities + should use the input dtype. + + Returns: + The prediction weights and probabilities followed by the target + weights and probabilities. + + Raises: + ValueError: if the configured kernel type is unsupported. + """ if self.kernel_type == "gaussian": - pred_weight, pred_probability = self.parzen_windowing_gaussian(pred) - target_weight, target_probability = self.parzen_windowing_gaussian(target) + pred_weight, pred_probability = self.parzen_windowing_gaussian(pred, restore_input_dtype) + target_weight, target_probability = self.parzen_windowing_gaussian(target, restore_input_dtype) elif self.kernel_type == "b-spline": # a third order BSpline kernel is used for the pred image intensity PDF. pred_weight, pred_probability = self.parzen_windowing_b_spline(pred, order=3) @@ -320,12 +335,16 @@ def parzen_windowing_b_spline(self, img: torch.Tensor, order: int) -> tuple[torc probability = torch.mean(weight, dim=-2, keepdim=True) # (batch, 1, num_bins) return weight, probability - def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def parzen_windowing_gaussian( + self, img: torch.Tensor, restore_input_dtype: bool = True + ) -> tuple[torch.Tensor, torch.Tensor]: """ Parzen windowing with gaussian kernel (adapted from DeepReg implementation) Note: the input is expected to range between 0 and 1 Args: img: the shape should be B[NDHW]. + restore_input_dtype: whether weights and probabilities should use + the input dtype. Returns: A tuple containing per-sample Gaussian bin weights and the @@ -350,7 +369,7 @@ def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, to weight = torch.exp(-preterm * (img - bin_centers) ** 2) # (batch, num_sample, num_bin) weight = weight / torch.sum(weight, dim=-1, keepdim=True) # (batch, num_sample, num_bin) probability = torch.mean(weight, dim=-2, keepdim=True) # (batch, 1, num_bin) - if output_dtype in (torch.float16, torch.bfloat16): + if restore_input_dtype and output_dtype in (torch.float16, torch.bfloat16): weight = weight.to(dtype=output_dtype) probability = probability.to(dtype=output_dtype) return weight, probability @@ -370,14 +389,16 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: """ if target.shape != pred.shape: raise ValueError(f"ground truth has differing shape ({target.shape}) from pred ({pred.shape})") - wa, pa, wb, pb = self.parzen_windowing(pred, target) # (batch, num_sample, num_bin), (batch, 1, num_bin) + wa, pa, wb, pb = self.parzen_windowing( + pred, target, restore_input_dtype=False + ) # (batch, num_sample, num_bin), (batch, 1, num_bin) # A half-precision matrix product can overflow while accumulating the # unnormalized joint histogram. Eager execution disables autocast for # this operation. TorchScript cannot compile a dynamic autocast device, # so it computes the normalized histogram directly by scaling both # operands by sqrt(N). - output_dtype = wa.dtype + output_dtype = pred.dtype if self.kernel_type == "gaussian" else wa.dtype compute_dtype = torch.float32 if wa.dtype in (torch.float16, torch.bfloat16) else wa.dtype wa = wa.to(dtype=compute_dtype) wb = wb.to(wa) 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 d7ebd9b90c..a80c330834 100644 --- a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py +++ b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py @@ -256,6 +256,29 @@ def test_half_precision_nonconstant_images_match_float32(self, dtype): self.assertIsNotNone(pred.grad) self.assertTrue(torch.isfinite(pred.grad).all()) + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_half_precision_weak_mutual_information_matches_float32(self, dtype): + """Verify weak reduced-precision mutual information tracks float32. + + Args: + dtype: reduced-precision floating-point dtype to test. + """ + generator = torch.Generator().manual_seed(19) + pred_float = torch.rand((1, 1, 4096), generator=generator) + target_float = torch.rand((1, 1, 4096), generator=generator) + loss = GlobalMutualInformationLoss(kernel_type="gaussian", num_bins=8) + expected = loss(pred_float, target_float) + pred = pred_float.to(dtype=dtype).requires_grad_() + target = target_float.to(dtype=dtype) + + result = loss(pred, target) + + self.assertEqual(result.dtype, dtype) + torch.testing.assert_close(result.float(), expected, rtol=1e-2, atol=1e-6) + result.backward() + self.assertIsNotNone(pred.grad) + self.assertTrue(torch.isfinite(pred.grad).all()) + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) def test_half_precision_large_constant_volume_is_finite(self, dtype): """Verify reduced-precision loss and gradients remain finite.""" From ceb19147efccc3ffeb042cb5e30d2a21dd59b544 Mon Sep 17 00:00:00 2001 From: kyinhub Date: Wed, 29 Jul 2026 22:34:09 -0700 Subject: [PATCH 4/5] docs(losses): fix Gaussian Parzen docstring Signed-off-by: kyinhub --- monai/losses/image_dissimilarity.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 0f4d658b61..619e00f170 100644 --- a/monai/losses/image_dissimilarity.py +++ b/monai/losses/image_dissimilarity.py @@ -338,9 +338,10 @@ def parzen_windowing_b_spline(self, img: torch.Tensor, order: int) -> tuple[torc def parzen_windowing_gaussian( self, img: torch.Tensor, restore_input_dtype: bool = True ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Parzen windowing with gaussian kernel (adapted from DeepReg implementation) - Note: the input is expected to range between 0 and 1 + """Apply Gaussian Parzen windowing adapted from DeepReg. + + The input is expected to range between 0 and 1. + Args: img: the shape should be B[NDHW]. restore_input_dtype: whether weights and probabilities should use From 3dfbfe95cb3c7d45152c91ca01700aa93e0a8606 Mon Sep 17 00:00:00 2001 From: kyinhub Date: Wed, 29 Jul 2026 23:06:53 -0700 Subject: [PATCH 5/5] fix(losses): validate Gaussian MI parameters Signed-off-by: kyinhub --- monai/losses/image_dissimilarity.py | 11 ++++++- .../test_global_mutual_information_loss.py | 33 +++++++++++++++++-- 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 619e00f170..1282df32ee 100644 --- a/monai/losses/image_dissimilarity.py +++ b/monai/losses/image_dissimilarity.py @@ -226,13 +226,22 @@ 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 the Gaussian kernel + has fewer than two bins or a non-finite, non-positive ``sigma_ratio``. """ super().__init__(reduction=LossReduction(reduction).value) + self.kernel_type = look_up_option(kernel_type, ["gaussian", "b-spline"]) if num_bins <= 0: raise ValueError(f"num_bins must > 0, got {num_bins}") + if self.kernel_type == "gaussian": + if num_bins < 2: + raise ValueError(f"Gaussian kernel requires num_bins >= 2, got {num_bins}") + if not math.isfinite(sigma_ratio) or sigma_ratio <= 0.0: + raise ValueError(f"Gaussian kernel requires a finite, positive sigma_ratio, got {sigma_ratio}") 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"]) 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. 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 a80c330834..e15db6cca2 100644 --- a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py +++ b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py @@ -163,13 +163,38 @@ def test_ill_opts(self, num_bins, reduction, expected_exception, expected_messag with self.assertRaisesRegex(expected_exception, expected_message): GlobalMutualInformationLoss(num_bins=num_bins, reduction=reduction)(pred, target) + @parameterized.expand( + [ + (1, 0.5, "num_bins >= 2"), + (23, 0.0, "finite, positive sigma_ratio"), + (23, -0.5, "finite, positive sigma_ratio"), + (23, float("nan"), "finite, positive sigma_ratio"), + (23, float("inf"), "finite, positive sigma_ratio"), + (23, float("-inf"), "finite, positive sigma_ratio"), + ] + ) + def test_ill_gaussian_parameters(self, num_bins, sigma_ratio, expected_message): + """Verify invalid Gaussian parameters fail during construction. + + Args: + num_bins: number of histogram bins to test. + sigma_ratio: Gaussian kernel width ratio to test. + expected_message: text expected in the validation error. + """ + with self.assertRaisesRegex(ValueError, expected_message): + GlobalMutualInformationLoss(kernel_type="gaussian", num_bins=num_bins, sigma_ratio=sigma_ratio) + class TestGlobalMutualInformationLossHalfPrecision(unittest.TestCase): """Test stable Gaussian mutual information in reduced-precision modes.""" @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) def test_half_precision_gaussian_weights_with_many_bins_are_finite(self, dtype): - """Verify many-bin Parzen outputs remain finite and preserve metadata.""" + """Verify many-bin Parzen outputs remain finite and preserve metadata. + + Args: + dtype: reduced-precision floating-point dtype to test. + """ image = torch.zeros((1, 1, 2), dtype=dtype) loss = GlobalMutualInformationLoss(kernel_type="gaussian", num_bins=256) @@ -281,7 +306,11 @@ def test_half_precision_weak_mutual_information_matches_float32(self, dtype): @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) def test_half_precision_large_constant_volume_is_finite(self, dtype): - """Verify reduced-precision loss and gradients remain finite.""" + """Verify reduced-precision loss and gradients remain finite. + + Args: + dtype: reduced-precision floating-point dtype to test. + """ pred = torch.zeros((1, 1, 48, 48, 48), dtype=dtype, requires_grad=True) target = torch.zeros_like(pred) loss = GlobalMutualInformationLoss(kernel_type="gaussian")