fix(losses): avoid half precision mutual information overflow

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

Author: yinkev

monai/losses/image_dissimilarity.py

@@ -317,15 +317,22 @@ def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, to
         Args:
             img: the shape should be B[NDHW].
         """
-        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:
@@ -340,16 +347,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)

tests/losses/image_dissimilarity/test_global_mutual_information_loss.py

@@ -164,6 +164,48 @@ 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):
+    def test_half_precision_gaussian_weights_with_many_bins_are_finite(self):
+        image = torch.zeros((1, 1, 2), dtype=torch.float16)
+        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())
+
+    def test_half_precision_large_constant_volume_is_finite(self):
+        pred = torch.zeros((1, 1, 48, 48, 48), dtype=torch.float16, 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):
+        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."""