Prevent half-precision mutual information overflow

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)
@@ -316,40 +326,83 @@ 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
+        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 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:
         """
         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)
+        # 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)
+        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)

tests/losses/image_dissimilarity/test_global_mutual_information_loss.py

@@ -164,6 +164,165 @@ 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)
+        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):
+        """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_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):
         """Verify gaussian kernel registers preterm and bin_centers as non-trainable, non-persistent buffers."""