fix(losses): strengthen reduced-precision MI contracts

Signed-off-by: kyinhub <kevinpyin@gmail.com>

Author: yinkev

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

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())