diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 195ac32b1f..1282df32ee 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 @@ -224,32 +226,64 @@ 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. 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) 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) @@ -310,22 +344,44 @@ 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]: - """ - Parzen windowing with gaussian kernel (adapted from DeepReg implementation) - Note: the input is expected to range between 0 and 1 + def parzen_windowing_gaussian( + self, img: torch.Tensor, restore_input_dtype: bool = True + ) -> tuple[torch.Tensor, torch.Tensor]: + """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 + the input dtype. + + 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 + 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.") - 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) + 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) probability = torch.mean(weight, dim=-2, keepdim=True) # (batch, 1, num_bin) + 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 def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: @@ -333,23 +389,51 @@ 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"]. + 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) - - 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) + 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 = 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) + pa = pa.to(dtype=compute_dtype) + pb = pb.to(pa) + 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: - 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..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,6 +163,217 @@ 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. + + 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) + + 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) + 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_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. + + 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") + + 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_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) + 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)) + self.assertEqual(result.dtype, pred.dtype) + 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):