perf(viscy-utils): bf16-precision SSIM helper for Hopper FCMAE training

packages/viscy-utils/src/viscy_utils/evaluation/metrics.py

@@ -1,12 +1,12 @@
 """Metrics for model evaluation"""
 
+from math import prod
 from typing import Sequence, Union
 from warnings import warn
 
 import numpy as np
 import torch
 import torch.nn.functional as F
-from monai.metrics.regression import compute_ssim_and_cs
 from scipy.optimize import linear_sum_assignment
 from skimage.measure import label, regionprops
 from torchmetrics.detection.mean_ap import MeanAveragePrecision
@@ -42,10 +42,7 @@ def VOI_metric(target, prediction):
     im_intersection = np.logical_and(im_pred_mask, im_targ_mask)
     im_inters_informed = im_intersection * im_targ_mask * im_pred_mask
 
-    marg_intr = (
-        np.histogramdd(np.ravel(im_inters_informed), bins=256)[0]
-        / im_inters_informed.size
-    )
+    marg_intr = np.histogramdd(np.ravel(im_inters_informed), bins=256)[0] / im_inters_informed.size
     marg_intr = list(filter(lambda p: p > 0, np.ravel(marg_intr)))
     entropy_intr = -np.sum(np.multiply(marg_intr, np.log2(marg_intr)))
 
@@ -128,9 +125,7 @@ def labels_to_masks(labels: torch.ShortTensor) -> torch.BoolTensor:
         raise ValueError(f"Labels must be 2D, got shape {labels.shape}.")
     segments = torch.unique(labels)
     n_instances = segments.numel() - 1
-    masks = torch.zeros(
-        (n_instances, *labels.shape), dtype=torch.bool, device=labels.device
-    )
+    masks = torch.zeros((n_instances, *labels.shape), dtype=torch.bool, device=labels.device)
     # TODO: optimize this?
     for s, segment in enumerate(segments):
         # start from label value 1, i.e. skip background label
@@ -150,13 +145,9 @@ def labels_to_detection(labels: torch.ShortTensor) -> dict[str, torch.Tensor]:
     return {
         "boxes": boxes,
         # dummy confidence scores
-        "scores": torch.ones(
-            (boxes.shape[0],), dtype=torch.float32, device=boxes.device
-        ),
+        "scores": torch.ones((boxes.shape[0],), dtype=torch.float32, device=boxes.device),
         # dummy class labels
-        "labels": torch.zeros(
-            (boxes.shape[0],), dtype=torch.uint8, device=boxes.device
-        ),
+        "labels": torch.zeros((boxes.shape[0],), dtype=torch.uint8, device=boxes.device),
         "masks": masks,
     }
 
@@ -172,18 +163,112 @@ def mean_average_precision(
         :py:class:`torchmetrics.detection.MeanAveragePrecision`
     :return dict[str, torch.Tensor]: COCO-style metrics
     """
-    defaults = dict(
-        iou_type="segm", box_format="xyxy", max_detection_thresholds=[1, 100, 10000]
-    )
+    defaults = dict(iou_type="segm", box_format="xyxy", max_detection_thresholds=[1, 100, 10000])
     if not kwargs:
         kwargs = {}
     map_metric = MeanAveragePrecision(**(defaults | kwargs))
-    map_metric.update(
-        [labels_to_detection(pred_labels)], [labels_to_detection(target_labels)]
-    )
+    map_metric.update([labels_to_detection(pred_labels)], [labels_to_detection(target_labels)])
     return map_metric.compute()
 
 
+def _compute_ssim_and_cs_bf16(
+    y_pred: torch.Tensor,
+    y: torch.Tensor,
+    kernel_size: Sequence[int],
+    data_range: float | torch.Tensor = 1.0,
+    k1: float = 0.01,
+    k2: float = 0.03,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute SSIM and contrast-sensitivity with bf16 convolutions.
+
+    Replaces monai's ``compute_ssim_and_cs`` (which unconditionally casts both
+    inputs to fp32 internally) with a precision-aware variant that runs the
+    five uniform-window mean convolutions in bf16 and promotes only the
+    variance subtractions and C1/C2-guarded divisions to fp32. Squared
+    products (``y * y``, ``y_pred * y_pred``, ``y_pred * y``) are computed in
+    fp32 before casting to bf16 for the conv input, which preserves precision
+    on the precision-sensitive squaring step at the cost of one extra cast
+    per squared term.
+
+    The conv accumulator is fp32 on CUDA tensor cores (sm >= 80) and on CPU,
+    so only the conv outputs (returned in bf16) lose precision relative to
+    monai's all-fp32 path.
+
+    Specialized to 3D uniform kernels — matches the single call site in
+    :func:`ssim_25d`. Gaussian-kernel, ``kernel_sigma``, and ``spatial_dims``
+    parameters from monai's signature are intentionally dropped.
+
+    Parameters
+    ----------
+    y_pred : torch.Tensor
+        Predicted batch with shape ``(B, C, D, H, W)``.
+    y : torch.Tensor
+        Target batch with the same shape as ``y_pred``.
+    kernel_size : Sequence[int]
+        Uniform 3D window size ``(D, H, W)``.
+    data_range : float or torch.Tensor, optional
+        Data range of the inputs; used to compute the C1, C2 stability
+        constants. Defaults to ``1.0``.
+    k1 : float, optional
+        Luminance stability constant. Defaults to ``0.01``.
+    k2 : float, optional
+        Contrast stability constant. Defaults to ``0.03``.
+
+    Returns
+    -------
+    ssim : torch.Tensor
+        Per-pixel SSIM map in fp32, shape ``(B, C, *reduced_spatial)``.
+    cs : torch.Tensor
+        Per-pixel contrast-sensitivity map in fp32, same shape as ``ssim``.
+    """
+    if y.shape != y_pred.shape:
+        raise ValueError(f"y_pred and y must have same shape, got {y_pred.shape} and {y.shape}.")
+
+    num_channels = y_pred.size(1)
+
+    # Build uniform kernel in fp32 then cast to bf16 once.
+    # The kernel is rebuilt per call (5x per ms_ssim_25d invocation). The
+    # tensor is small (~1.8k elements at default 15x11x11) so the per-step
+    # overhead is negligible relative to the conv cost. If profiling ever
+    # shows otherwise, cache by (num_channels, kernel_size, device, dtype).
+    kernel_fp32 = torch.ones((num_channels, 1, *kernel_size), device=y_pred.device, dtype=torch.float32) / float(
+        prod(kernel_size)
+    )
+    kernel_bf = kernel_fp32.to(torch.bfloat16)
+
+    # bf16 views of the simple (non-squared) inputs go straight from the
+    # caller's dtype — skips a round-trip via fp32 when the caller is
+    # already in bf16 (autocast). Squared products are computed in fp32
+    # first to preserve squaring precision, then cast to bf16 for the conv.
+    y_pred_bf = y_pred.to(torch.bfloat16)
+    y_bf = y.to(torch.bfloat16)
+    y_pred_fp32 = y_pred.float()
+    y_fp32 = y.float()
+    y_pred_sq_bf = (y_pred_fp32 * y_pred_fp32).to(torch.bfloat16)
+    y_sq_bf = (y_fp32 * y_fp32).to(torch.bfloat16)
+    y_pred_y_bf = (y_pred_fp32 * y_fp32).to(torch.bfloat16)
+
+    mu_x = F.conv3d(y_pred_bf, kernel_bf, groups=num_channels).float()
+    mu_y = F.conv3d(y_bf, kernel_bf, groups=num_channels).float()
+    mu_xx = F.conv3d(y_pred_sq_bf, kernel_bf, groups=num_channels).float()
+    mu_yy = F.conv3d(y_sq_bf, kernel_bf, groups=num_channels).float()
+    mu_xy = F.conv3d(y_pred_y_bf, kernel_bf, groups=num_channels).float()
+
+    # Stability constants in fp32 (data_range may be a 0-dim tensor; the
+    # multiplication promotes to fp32 since k1/k2 are Python floats).
+    c1 = (k1 * data_range) ** 2
+    c2 = (k2 * data_range) ** 2
+
+    sigma_x = mu_xx - mu_x * mu_x
+    sigma_y = mu_yy - mu_y * mu_y
+    sigma_xy = mu_xy - mu_x * mu_y
+
+    contrast_sensitivity = (2 * sigma_xy + c2) / (sigma_x + sigma_y + c2)
+    ssim_full = ((2 * mu_x * mu_y + c1) / (mu_x * mu_x + mu_y * mu_y + c1)) * contrast_sensitivity
+
+    return ssim_full, contrast_sensitivity
+
+
 def ssim_25d(
     preds: torch.Tensor,
     target: torch.Tensor,
@@ -202,20 +287,15 @@ def ssim_25d(
     :return Optional[torch.Tensor]: contrast sensitivity
     """
     if preds.ndim != 5:
-        raise ValueError(
-            f"Input shape must be (B, C, D, W, H), got input shape {preds.shape}"
-        )
+        raise ValueError(f"Input shape must be (B, C, D, W, H), got input shape {preds.shape}")
     depth = preds.shape[2]
     if depth > 15:
         warn(f"Input depth {depth} is potentially too large for 2.5D SSIM.")
-    ssim_img, cs_img = compute_ssim_and_cs(
+    ssim_img, cs_img = _compute_ssim_and_cs_bf16(
         preds,
         target,
-        3,
-        kernel_sigma=None,
         kernel_size=(depth, *in_plane_window_size),
         data_range=target.max(),
-        kernel_type="uniform",
     )
     # aggregate to one scalar per batch
     ssim = ssim_img.view(ssim_img.shape[0], -1).mean(1)
@@ -253,9 +333,7 @@ def ms_ssim_25d(
     base_min = 1e-4
     mcs_list = []
     for _ in range(len(betas)):
-        ssim, contrast_sensitivity = ssim_25d(
-            preds, target, in_plane_window_size, return_contrast_sensitivity=True
-        )
+        ssim, contrast_sensitivity = ssim_25d(preds, target, in_plane_window_size, return_contrast_sensitivity=True)
         if clamp:
             contrast_sensitivity = contrast_sensitivity.clamp(min=base_min)
         mcs_list.append(contrast_sensitivity)

packages/viscy-utils/src/viscy_utils/losses/mixed_loss.py

@@ -4,7 +4,6 @@
 for image reconstruction tasks, adapted from Zhao et al.
 """
 
-import torch
 import torch.nn.functional as F
 from torch import nn
 
@@ -40,7 +39,6 @@ def __init__(
         self.l2_alpha = l2_alpha
         self.ms_dssim_alpha = ms_dssim_alpha
 
-    @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float32)
     def forward(self, preds, target):
         """Compute the mixed reconstruction loss.
 

packages/viscy-utils/src/viscy_utils/losses/spotlight.py

@@ -158,7 +158,6 @@ def __init__(
         self.fg_threshold = fg_threshold
         self._warned_no_real_mask = False
 
-    @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float32)
     def forward(self, pred: Tensor, target: Tensor, fg_mask: Tensor | None = None) -> Tensor:
         """Compute the Spotlight loss.
 

packages/viscy-utils/tests/test_metrics.py

@@ -0,0 +1,169 @@
+"""Tests for the bf16-precision SSIM helper in viscy_utils.evaluation.metrics.
+
+Covers the multi-tier numerical contract:
+
+- per-pixel SSIM equivalence on random inputs (worst-case bf16 drift)
+- aggregate SSIM equivalence on random inputs (per-pixel noise averages out)
+- aggregate SSIM equivalence on correlated-pair inputs (closer to training)
+- gradient-flow correctness via cosine similarity and sign-flip fraction
+- output dtype invariance to input dtype
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from viscy_utils.evaluation.metrics import _compute_ssim_and_cs_bf16
+
+# monai is not a hard dep of viscy-utils — skip the suite if absent rather
+# than failing at import time.
+_monai_regression = pytest.importorskip("monai.metrics.regression")
+_monai_reference = _monai_regression.compute_ssim_and_cs
+
+# The helper unconditionally uses bf16 convs. CUDA bf16 conv works on
+# sm_80+ in tensor cores and falls back to software emulation on older
+# devices, but the equivalence-vs-monai-fp32 tolerances were measured on
+# Hopper — skip on hardware where bf16 emulation could push drift past
+# the configured rtol/atol.
+_skip_no_bf16 = pytest.mark.skipif(
+    not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
+    reason="CUDA + bf16 tensor-core support required",
+)
+
+# Representative iPSC SEC61B FCMAE batch shape.
+_BATCH = (2, 1, 15, 256, 256)
+_KERNEL = (15, 11, 11)
+
+
+def _ref(y_pred: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return _monai_reference(
+        y_pred,
+        y,
+        spatial_dims=3,
+        kernel_size=_KERNEL,
+        kernel_sigma=None,
+        kernel_type="uniform",
+        data_range=y.max(),
+    )
+
+
+def _bf16(y_pred: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return _compute_ssim_and_cs_bf16(
+        y_pred,
+        y,
+        kernel_size=_KERNEL,
+        data_range=y.max(),
+    )
+
+
+@_skip_no_bf16
+def test_ssim_helper_random_per_pixel_equivalence():
+    """Per-pixel SSIM on random inputs, worst-case bf16 drift tier.
+
+    Tolerance ≥2× margin over measured 0.0418 absolute drift.
+    """
+    torch.manual_seed(0)
+    y_pred = torch.rand(*_BATCH, device="cuda")
+    y = torch.rand(*_BATCH, device="cuda")
+
+    ssim_ref, cs_ref = _ref(y_pred, y)
+    ssim_helper, cs_helper = _bf16(y_pred, y)
+
+    torch.testing.assert_close(ssim_helper, ssim_ref, rtol=5e-2, atol=1e-1)
+    torch.testing.assert_close(cs_helper, cs_ref, rtol=5e-2, atol=1e-1)
+
+
+@_skip_no_bf16
+def test_ssim_helper_random_aggregate_equivalence():
+    """Aggregate SSIM (mean over pixels) on random inputs.
+
+    Per-pixel noise averages out across H×W ≈ 65k pixels.
+    Tolerance ≥25% margin over measured 0.00776 absolute drift.
+    """
+    torch.manual_seed(1)
+    y_pred = torch.rand(*_BATCH, device="cuda")
+    y = torch.rand(*_BATCH, device="cuda")
+
+    ssim_ref, _ = _ref(y_pred, y)
+    ssim_helper, _ = _bf16(y_pred, y)
+
+    agg_ref = ssim_ref.view(ssim_ref.shape[0], -1).mean(1)
+    agg_helper = ssim_helper.view(ssim_helper.shape[0], -1).mean(1)
+
+    torch.testing.assert_close(agg_helper, agg_ref, rtol=1e-2, atol=1e-2)
+
+
+@_skip_no_bf16
+def test_ssim_helper_correlated_pair_equivalence():
+    """Aggregate SSIM on a correlated pair (pred = target + small noise).
+
+    Closer to training-time inputs; SSIM lives near 1.0 so relative drift
+    is much smaller than on uncorrelated random data.
+    """
+    torch.manual_seed(2)
+    y = torch.rand(*_BATCH, device="cuda")
+    y_pred = y + 0.05 * torch.randn_like(y)
+
+    ssim_ref, _ = _ref(y_pred, y)
+    ssim_helper, _ = _bf16(y_pred, y)
+
+    agg_ref = ssim_ref.view(ssim_ref.shape[0], -1).mean(1)
+    agg_helper = ssim_helper.view(ssim_helper.shape[0], -1).mean(1)
+
+    torch.testing.assert_close(agg_helper, agg_ref, rtol=2e-3, atol=5e-3)
+
+
+@_skip_no_bf16
+def test_ssim_helper_gradient_flow():
+    """Gradient flow contract: finite, cosine-similar, low sign-flip rate.
+
+    Per-voxel sign equality is too brittle (~0.25% benign flips on
+    non-tiny gradients). Use cosine similarity + sign-flip fraction over
+    voxels above 10% of the reference grad max — relative threshold so
+    the assertion is scale-invariant regardless of loss magnitude.
+    """
+    torch.manual_seed(3)
+    y = torch.rand(*_BATCH, device="cuda")
+
+    y_pred_ref = (y + 0.05 * torch.randn_like(y)).detach().requires_grad_(True)
+    ssim_ref, _ = _ref(y_pred_ref, y)
+    (1 - ssim_ref).mean().backward()
+    grad_ref = y_pred_ref.grad
+
+    y_pred_helper = y_pred_ref.detach().clone().requires_grad_(True)
+    ssim_helper, _ = _bf16(y_pred_helper, y)
+    (1 - ssim_helper).mean().backward()
+    grad_helper = y_pred_helper.grad
+
+    assert grad_helper is not None
+    assert grad_helper.shape == grad_ref.shape
+    assert torch.isfinite(grad_helper).all()
+
+    cos_sim = F.cosine_similarity(
+        grad_helper.flatten().unsqueeze(0),
+        grad_ref.flatten().unsqueeze(0),
+        dim=1,
+    ).item()
+    assert cos_sim >= 0.99, f"cosine similarity {cos_sim:.4f} below 0.99"
+
+    # Relative threshold — observed |grad_ref| max is ~1.7e-6, so the
+    # earlier absolute 1e-3 threshold was vacuous. Anchor to 10% of the
+    # reference max so the assertion is scale-invariant and meaningful.
+    nontiny = grad_ref.abs() > 0.1 * grad_ref.abs().max()
+    assert nontiny.any(), "no non-tiny reference gradients to compare signs against"
+    flip_fraction = ((grad_helper.sign() != grad_ref.sign()) & nontiny).float().sum() / nontiny.float().sum()
+    assert flip_fraction.item() < 0.01, f"sign-flip fraction {flip_fraction.item():.4f} above 1%"
+
+
+@_skip_no_bf16
+@pytest.mark.parametrize("input_dtype", [torch.float32, torch.bfloat16, torch.float16])
+def test_ssim_helper_dtypes(input_dtype):
+    """Helper returns fp32 regardless of input dtype."""
+    torch.manual_seed(4)
+    y_pred = torch.rand(*_BATCH, device="cuda", dtype=input_dtype)
+    y = torch.rand(*_BATCH, device="cuda", dtype=input_dtype)
+
+    ssim_helper, cs_helper = _bf16(y_pred, y)
+
+    assert ssim_helper.dtype == torch.float32
+    assert cs_helper.dtype == torch.float32