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
@@ -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,46 +344,96 @@ 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:
         """
         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)