From 5384817f026c9b82d1df30d6a9a1de3ccd1222b6 Mon Sep 17 00:00:00 2001 From: Theo Barfoot Date: Mon, 31 Aug 2026 15:17:53 +0100 Subject: [PATCH 1/2] Add average calibration losses Signed-off-by: Theo Barfoot --- docs/source/losses.rst | 10 + monai/handlers/calibration.py | 4 +- monai/losses/__init__.py | 1 + monai/losses/calibration.py | 377 ++++++++++++++++++++++++++ monai/metrics/calibration.py | 8 +- tests/losses/test_calibration_loss.py | 278 +++++++++++++++++++ 6 files changed, 672 insertions(+), 6 deletions(-) create mode 100644 monai/losses/calibration.py create mode 100644 tests/losses/test_calibration_loss.py diff --git a/docs/source/losses.rst b/docs/source/losses.rst index a5be560324b..1dca79a1aca 100644 --- a/docs/source/losses.rst +++ b/docs/source/losses.rst @@ -22,6 +22,16 @@ Segmentation Losses .. autoclass:: dice :members: +`HardL1ACELoss` +~~~~~~~~~~~~~~~~ +.. autoclass:: HardL1ACELoss + :members: + +`SoftL1ACELoss` +~~~~~~~~~~~~~~~~ +.. autoclass:: SoftL1ACELoss + :members: + `MaskedDiceLoss` ~~~~~~~~~~~~~~~~ .. autoclass:: MaskedDiceLoss diff --git a/monai/handlers/calibration.py b/monai/handlers/calibration.py index addcc572303..b8365c9c680 100644 --- a/monai/handlers/calibration.py +++ b/monai/handlers/calibration.py @@ -59,8 +59,8 @@ class CalibrationError(IgniteMetricHandler): - Guo, C., et al. "On Calibration of Modern Neural Networks." ICML 2017. https://proceedings.mlr.press/v70/guo17a.html - Barfoot, T., et al. "Average Calibration Losses for Reliable Uncertainty in - Medical Image Segmentation." arXiv:2506.03942v3, 2025. - https://arxiv.org/abs/2506.03942v3 + Medical Image Segmentation." IEEE Transactions on Medical Imaging, 2026. + https://doi.org/10.1109/TMI.2026.3673118 See Also: - :py:class:`~monai.metrics.CalibrationErrorMetric`: The underlying metric class. diff --git a/monai/losses/__init__.py b/monai/losses/__init__.py index 9f35e5f0750..c6bc5e3810f 100644 --- a/monai/losses/__init__.py +++ b/monai/losses/__init__.py @@ -15,6 +15,7 @@ from .aucm_loss import AUCMLoss from .barlow_twins import BarlowTwinsLoss from .boundary_loss import BoundaryLoss +from .calibration import HardL1ACELoss, SoftL1ACELoss from .cldice import SoftclDiceLoss, SoftDiceclDiceLoss from .contrastive import ContrastiveLoss from .deform import BendingEnergyLoss, DiffusionLoss diff --git a/monai/losses/calibration.py b/monai/losses/calibration.py new file mode 100644 index 00000000000..5e8b5b0751a --- /dev/null +++ b/monai/losses/calibration.py @@ -0,0 +1,377 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import warnings +from collections.abc import Callable, Sequence + +import torch +from torch.nn.modules.loss import _Loss + +from monai.networks import one_hot +from monai.utils import LossReduction + +__all__ = ["HardL1ACELoss", "SoftL1ACELoss"] + + +def _accumulation_dtype(input: torch.Tensor) -> torch.dtype: + return torch.float32 if input.dtype in (torch.float16, torch.bfloat16) else input.dtype + + +def _hard_binned_calibration( + input: torch.Tensor, target: torch.Tensor, num_bins: int, right: bool +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return hard-binned prediction sums, target sums, and counts.""" + work_dtype = _accumulation_dtype(input) + input_flat = input.flatten(start_dim=2).to(dtype=work_dtype).contiguous() + target_flat = target.detach().flatten(start_dim=2).to(dtype=work_dtype) + + # Match calibration_binning's established boundaries, including its epsilon-expanded upper edge. + # Spell out float32 epsilon because torch.finfo is not supported by TorchScript. + float32_eps = 1.1920928955078125e-7 + boundaries = torch.linspace(0.0, 1.0 + float32_eps, num_bins + 1, dtype=work_dtype, device=input.device) + bin_idx = torch.bucketize(input_flat, boundaries[1:], right=right).clamp(max=num_bins - 1) + counts = torch.zeros(input_flat.shape[0], input_flat.shape[1], num_bins, dtype=work_dtype, device=input.device) + counts = counts.scatter_add(2, bin_idx, torch.ones_like(input_flat)) + sum_p = torch.zeros_like(counts).scatter_add(2, bin_idx, input_flat) + sum_target = torch.zeros_like(counts).scatter_add(2, bin_idx, target_flat) + return sum_p, sum_target, counts + + +def _soft_binned_calibration( + input: torch.Tensor, target: torch.Tensor, num_bins: int, right: bool +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return soft-binned prediction sums, target sums, and effective counts.""" + work_dtype = _accumulation_dtype(input) + input_flat = input.flatten(start_dim=2).to(dtype=work_dtype).contiguous() + target_flat = target.detach().flatten(start_dim=2).to(dtype=work_dtype) + + # Spell out float32 epsilon because torch.finfo is not supported by TorchScript. + float32_eps = 1.1920928955078125e-7 + half_boundaries = torch.linspace(0.0, 1.0 + float32_eps, 2 * num_bins + 1, dtype=work_dtype, device=input.device) + centers = half_boundaries[1::2].contiguous() + insertion_idx = torch.bucketize(input_flat, centers, right=right) + left_idx = (insertion_idx - 1).clamp(min=0, max=num_bins - 1) + right_idx = insertion_idx.clamp(max=num_bins - 1) + + left_centers = centers[left_idx] + right_centers = centers[right_idx] + distinct = left_idx != right_idx + distance = (right_centers - left_centers).clamp_min(float32_eps) + right_weight = torch.where(distinct, (input_flat - left_centers) / distance, torch.zeros_like(input_flat)) + left_weight = 1.0 - right_weight + + counts = torch.zeros(input_flat.shape[0], input_flat.shape[1], num_bins, dtype=work_dtype, device=input.device) + counts = counts.scatter_add(2, left_idx, left_weight).scatter_add(2, right_idx, right_weight) + sum_p = torch.zeros_like(counts) + sum_p = sum_p.scatter_add(2, left_idx, left_weight * input_flat) + sum_p = sum_p.scatter_add(2, right_idx, right_weight * input_flat) + sum_target = torch.zeros_like(counts) + sum_target = sum_target.scatter_add(2, left_idx, left_weight * target_flat) + sum_target = sum_target.scatter_add(2, right_idx, right_weight * target_flat) + return sum_p, sum_target, counts + + +class _L1ACELoss(_Loss): + """Shared input handling and reduction for marginal L1 ACE losses.""" + + def __init__( + self, + num_bins: int, + include_background: bool, + to_onehot_y: bool, + sigmoid: bool, + softmax: bool, + other_act: Callable | None, + reduction: LossReduction | str, + weight: Sequence[float] | float | int | torch.Tensor | None, + right: bool, + ignore_empty_classes: bool, + ) -> None: + super().__init__(reduction=LossReduction(reduction).value) + if num_bins < 1: + raise ValueError(f"num_bins must be >= 1, got {num_bins}.") + if other_act is not None and not callable(other_act): + raise TypeError(f"other_act must be None or callable but is {type(other_act).__name__}.") + if int(sigmoid) + int(softmax) + int(other_act is not None) > 1: + raise ValueError("Incompatible values: more than 1 of [sigmoid=True, softmax=True, other_act is not None].") + + class_weight = torch.as_tensor(weight) if weight is not None else None + if class_weight is not None: + if class_weight.ndim > 1: + raise ValueError("weight must be a scalar or a one-dimensional sequence.") + if torch.any(class_weight < 0): + raise ValueError("the value/values of the `weight` should be no less than 0.") + + self.num_bins = num_bins + self.include_background = include_background + self.to_onehot_y = to_onehot_y + self.sigmoid = sigmoid + self.softmax = softmax + self.other_act = other_act + self.right = right + self.ignore_empty_classes = ignore_empty_classes + self.register_buffer("class_weight", class_weight) + self.class_weight: None | torch.Tensor + + def _prepare_input(self, input: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if input.ndim < 3: + raise ValueError(f"input must have shape (B, C, spatial...), got ndim={input.ndim}.") + if not input.is_floating_point(): + raise TypeError(f"input must be a floating point tensor, got {input.dtype}.") + + if self.sigmoid: + input = torch.sigmoid(input) + + n_pred_ch = input.shape[1] + if self.softmax: + if n_pred_ch == 1: + warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=3) + else: + input = torch.softmax(input, 1) + if self.other_act is not None: + input = self.other_act(input) + + if self.to_onehot_y: + if n_pred_ch == 1: + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=3) + else: + target = one_hot(target, num_classes=n_pred_ch) + + if not self.include_background: + if n_pred_ch == 1: + warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=3) + else: + input = input[:, 1:] + target = target[:, 1:] + + if target.shape != input.shape: + raise AssertionError(f"ground truth has different shape ({target.shape}) from input ({input.shape})") + return input, target + + def _reduce(self, per_class_loss: torch.Tensor, valid_classes: torch.Tensor, input: torch.Tensor) -> torch.Tensor: + num_classes = per_class_loss.shape[1] + if self.class_weight is not None: + if self.class_weight.ndim == 0: + class_weight = self.class_weight.expand(num_classes) + elif self.class_weight.shape[0] == num_classes: + class_weight = self.class_weight + else: + raise ValueError( + "The length of the `weight` sequence should be the same as the number of classes. " + "If `include_background=False`, the weight should not include the background category class 0." + ) + per_class_loss = per_class_loss * class_weight.to(per_class_loss) + + per_class_loss = per_class_loss * valid_classes.to(dtype=per_class_loss.dtype) + if self.reduction == LossReduction.MEAN.value: + return per_class_loss.mean() + if self.reduction == LossReduction.SUM.value: + return per_class_loss.sum() + if self.reduction == LossReduction.NONE.value: + broadcast_shape = list(per_class_loss.shape) + [1] * (input.ndim - 2) + return per_class_loss.view(broadcast_shape) + raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].') + + def _finish( + self, + sum_p: torch.Tensor, + sum_target: torch.Tensor, + counts: torch.Tensor, + valid_bins: torch.Tensor, + input: torch.Tensor, + target: torch.Tensor, + ) -> torch.Tensor: + safe_counts = torch.where(valid_bins, counts, torch.ones_like(counts)) + gap = torch.abs(sum_p / safe_counts - sum_target / safe_counts) + valid_bins_float = valid_bins.to(dtype=gap.dtype) + valid_bin_count = valid_bins_float.sum(dim=-1) + per_class_loss = (gap * valid_bins_float).sum(dim=-1) / valid_bin_count.clamp_min(1) + valid_classes = valid_bin_count > 0 + if self.ignore_empty_classes: + valid_classes = valid_classes & (target.flatten(start_dim=2).sum(dim=-1) > 0) + return self._reduce(per_class_loss, valid_classes, input) + + +class HardL1ACELoss(_L1ACELoss): + """ + Compute hard-binned marginal L1 Average Calibration Error (ACE) loss. + + The loss measures calibration independently for every image and class. Predicted probabilities are assigned to + hard bins, the absolute difference between mean probability and mean binary target is computed in each occupied + bin, and those differences are averaged equally. Unlike Expected Calibration Error, occupied bins are not + weighted by voxel count. Hard assignments are discrete, while the mean probability within a fixed assignment is + differentiable. + + Input must have shape ``(B, C, spatial...)``. It is interpreted as probabilities unless ``sigmoid``, ``softmax``, + or ``other_act`` is selected. Targets may have the same one-hot shape or be label maps of shape + ``(B, 1, spatial...)`` when ``to_onehot_y=True``. Ignored empty target classes contribute zero before reduction, + matching the reference implementation; with ``reduction="none"`` they are returned as zero. Class weights are + applied before reduction. + The unreduced shape is ``(B, C, 1, ..., 1)`` after optional background removal. + + Finite-bin calibration estimates depend on the bin count and data distribution. This loss is intended as an + auxiliary objective and may trade segmentation accuracy against calibration quality; tune its coefficient and + ``num_bins`` on validation data. + + Args: + num_bins: number of equally spaced bins. Defaults to 20. + include_background: whether channel 0 contributes. Defaults to ``True``. + to_onehot_y: convert a single-channel label map to one-hot targets. Defaults to ``False``. + sigmoid: apply sigmoid to input. + softmax: apply channel-wise softmax to input. + other_act: optional callable activation. Only one activation option may be used. + reduction: one of ``"none"``, ``"mean"``, or ``"sum"``. + weight: scalar or one non-negative value per included class. + right: hard-bin boundary inclusion rule, matching :py:func:`monai.metrics.calibration_binning`. + ignore_empty_classes: set the loss for classes absent from an image to zero before reduction. Defaults to + ``True``. + + See Also: + - :py:func:`monai.metrics.calibration_binning`: The corresponding calibration bin statistics. + - :py:class:`monai.metrics.CalibrationErrorMetric`: Evaluation metrics computed from those statistics. + + References: + - Barfoot et al., "Average Calibration Losses for Reliable Uncertainty in Medical Image Segmentation," + IEEE Transactions on Medical Imaging, 2026. https://doi.org/10.1109/TMI.2026.3673118 + - Barfoot et al., MICCAI 2024. https://papers.miccai.org/miccai-2024/091-Paper3075.html + - Guo et al., "On Calibration of Modern Neural Networks," ICML 2017. + + Example: + >>> import torch + >>> from monai.losses import DiceCELoss, HardL1ACELoss + >>> logits, labels = torch.randn(2, 3, 16, 16), torch.randint(0, 3, (2, 1, 16, 16)) + >>> segmentation = DiceCELoss(to_onehot_y=True, softmax=True) + >>> calibration = HardL1ACELoss(to_onehot_y=True, softmax=True) + >>> loss = segmentation(logits, labels) + 0.1 * calibration(logits, labels) + """ + + def __init__( + self, + num_bins: int = 20, + include_background: bool = True, + to_onehot_y: bool = False, + sigmoid: bool = False, + softmax: bool = False, + other_act: Callable | None = None, + reduction: LossReduction | str = LossReduction.MEAN, + weight: Sequence[float] | float | int | torch.Tensor | None = None, + right: bool = False, + ignore_empty_classes: bool = True, + ) -> None: + super().__init__( + num_bins, + include_background, + to_onehot_y, + sigmoid, + softmax, + other_act, + reduction, + weight, + right, + ignore_empty_classes, + ) + + def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Compute the loss between ``input`` and ``target``.""" + input, target = self._prepare_input(input, target) + sum_p, sum_target, counts = _hard_binned_calibration(input, target, self.num_bins, self.right) + return self._finish(sum_p, sum_target, counts, counts > 0, input, target) + + +class SoftL1ACELoss(_L1ACELoss): + """ + Compute soft-binned marginal L1 Average Calibration Error (ACE) loss. + + This loss has the same per-image, per-class and occupied-bin averaging semantics as :class:`HardL1ACELoss`, but + linearly interpolates each probability between its two neighboring bin centers. The implementation stores only + two indices and weights per probability, keeping memory complexity ``O(B*C*N)`` rather than materializing an + ``(B, C, N, num_bins)`` tensor. Soft binning provides a smoother training signal across bin transitions. + + Input must have shape ``(B, C, spatial...)`` and contain probabilities unless an activation is enabled. Targets + may be matching one-hot tensors or single-channel label maps when ``to_onehot_y=True``. Bins whose effective + weight is below ``empty_weight`` are ignored. Empty target classes can contribute zero before reduction, and + ``reduction="none"`` returns ``(B, C, 1, ..., 1)`` with ignored entries set to zero. + + Finite-bin estimates and soft assignments depend on ``num_bins``, ``empty_weight``, and the sample distribution. + Use this loss as a validated auxiliary objective: improved calibration can coincide with lower segmentation + performance. + + Args: + num_bins: number of equally spaced bin centers. Defaults to 20. + include_background: whether channel 0 contributes. Defaults to ``True``. + to_onehot_y: convert a single-channel label map to one-hot targets. Defaults to ``False``. + sigmoid: apply sigmoid to input. + softmax: apply channel-wise softmax to input. + other_act: optional callable activation. Only one activation option may be used. + reduction: one of ``"none"``, ``"mean"``, or ``"sum"``. + weight: scalar or one non-negative value per included class. + empty_weight: minimum effective bin weight. Empty bins are always ignored. Defaults to 0.01. + right: boundary rule used when a probability equals a bin center. + ignore_empty_classes: set the loss for classes absent from an image to zero before reduction. Defaults to + ``True``. + + See Also: + - :py:func:`monai.metrics.calibration_binning`: Hard-binned evaluation statistics for reliability diagrams. + - :py:class:`monai.metrics.CalibrationErrorMetric`: Evaluation metrics computed from those statistics. + + References: + - Barfoot et al., "Average Calibration Losses for Reliable Uncertainty in Medical Image Segmentation," + IEEE Transactions on Medical Imaging, 2026. https://doi.org/10.1109/TMI.2026.3673118 + - Barfoot et al., MICCAI 2024. https://papers.miccai.org/miccai-2024/091-Paper3075.html + - Guo et al., "On Calibration of Modern Neural Networks," ICML 2017. + + Example: + >>> import torch + >>> from monai.losses import DiceCELoss, SoftL1ACELoss + >>> logits, labels = torch.randn(2, 3, 16, 16), torch.randint(0, 3, (2, 1, 16, 16)) + >>> segmentation = DiceCELoss(to_onehot_y=True, softmax=True) + >>> calibration = SoftL1ACELoss(to_onehot_y=True, softmax=True) + >>> loss = segmentation(logits, labels) + 0.1 * calibration(logits, labels) + """ + + def __init__( + self, + num_bins: int = 20, + include_background: bool = True, + to_onehot_y: bool = False, + sigmoid: bool = False, + softmax: bool = False, + other_act: Callable | None = None, + reduction: LossReduction | str = LossReduction.MEAN, + weight: Sequence[float] | float | int | torch.Tensor | None = None, + empty_weight: float = 0.01, + right: bool = False, + ignore_empty_classes: bool = True, + ) -> None: + if empty_weight < 0: + raise ValueError(f"empty_weight must be >= 0, got {empty_weight}.") + super().__init__( + num_bins, + include_background, + to_onehot_y, + sigmoid, + softmax, + other_act, + reduction, + weight, + right, + ignore_empty_classes, + ) + self.empty_weight = float(empty_weight) + + def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Compute the loss between ``input`` and ``target``.""" + input, target = self._prepare_input(input, target) + sum_p, sum_target, counts = _soft_binned_calibration(input, target, self.num_bins, self.right) + valid_bins = (counts > 0) & (counts >= self.empty_weight) + return self._finish(sum_p, sum_target, counts, valid_bins, input, target) diff --git a/monai/metrics/calibration.py b/monai/metrics/calibration.py index 4231eeef993..a94c04037b8 100644 --- a/monai/metrics/calibration.py +++ b/monai/metrics/calibration.py @@ -69,8 +69,8 @@ def calibration_binning( - Guo, C., et al. "On Calibration of Modern Neural Networks." ICML 2017. https://proceedings.mlr.press/v70/guo17a.html - Barfoot, T., et al. "Average Calibration Losses for Reliable Uncertainty in - Medical Image Segmentation." arXiv:2506.03942v3, 2025. - https://arxiv.org/abs/2506.03942v3 + Medical Image Segmentation." IEEE Transactions on Medical Imaging, 2026. + https://doi.org/10.1109/TMI.2026.3673118 Note: This function uses nested loops over batch and channel dimensions for binning operations. @@ -211,8 +211,8 @@ class CalibrationErrorMetric(CumulativeIterationMetric): - Guo, C., et al. "On Calibration of Modern Neural Networks." ICML 2017. https://proceedings.mlr.press/v70/guo17a.html - Barfoot, T., et al. "Average Calibration Losses for Reliable Uncertainty in - Medical Image Segmentation." arXiv:2506.03942v3, 2025. - https://arxiv.org/abs/2506.03942v3 + Medical Image Segmentation." IEEE Transactions on Medical Imaging, 2026. + https://doi.org/10.1109/TMI.2026.3673118 See Also: - :py:class:`monai.handlers.CalibrationError`: Ignite handler wrapper for this metric. diff --git a/tests/losses/test_calibration_loss.py b/tests/losses/test_calibration_loss.py new file mode 100644 index 00000000000..d71775268f4 --- /dev/null +++ b/tests/losses/test_calibration_loss.py @@ -0,0 +1,278 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest +import warnings + +import torch + +from monai.losses import HardL1ACELoss, SoftL1ACELoss +from monai.losses.calibration import _hard_binned_calibration, _soft_binned_calibration +from monai.metrics import calibration_binning +from tests.test_utils import skip_if_no_cuda +from tests.test_utils import test_script_save as check_script_save + + +class TestCalibrationBinningHelpers(unittest.TestCase): + def test_hard_matches_metric_at_boundaries(self): + prediction = torch.tensor([[[0.0, 0.2, 0.4, 0.6, 0.8, 1.0]]]) + target = torch.tensor([[[0.0, 0.0, 1.0, 1.0, 1.0, 1.0]]]) + for right in (False, True): + with self.subTest(right=right): + sum_p, sum_target, counts = _hard_binned_calibration(prediction, target, 5, right) + metric_p, metric_target, metric_counts = calibration_binning(prediction, target, 5, right) + valid = counts > 0 + torch.testing.assert_close(counts, metric_counts) + torch.testing.assert_close(sum_p[valid] / counts[valid], metric_p[valid]) + torch.testing.assert_close(sum_target[valid] / counts[valid], metric_target[valid]) + + def test_soft_exact_interpolation(self): + prediction = torch.tensor([[[0.0, 0.25, 0.5, 0.75, 1.0]]], dtype=torch.float64) + target = torch.tensor([[[0.0, 0.0, 1.0, 1.0, 1.0]]], dtype=torch.float64) + sum_p, sum_target, counts = _soft_binned_calibration(prediction, target, 2, False) + torch.testing.assert_close(counts, torch.tensor([[[2.5, 2.5]]], dtype=torch.float64)) + torch.testing.assert_close(sum_p, torch.tensor([[[0.5, 2.0]]], dtype=torch.float64), rtol=0, atol=3e-7) + torch.testing.assert_close(sum_target, torch.tensor([[[0.5, 2.5]]], dtype=torch.float64), rtol=0, atol=3e-7) + + def test_soft_boundary_rule_and_continuity(self): + target = torch.tensor([[[1.0]]], dtype=torch.float64) + at_center = torch.tensor([[[0.3750000447034836]]], dtype=torch.float64) + left = at_center - 1e-7 + right = at_center + 1e-7 + for boundary_rule in (False, True): + with self.subTest(right=boundary_rule): + center_stats = _soft_binned_calibration(at_center, target, 4, boundary_rule) + left_stats = _soft_binned_calibration(left, target, 4, boundary_rule) + right_stats = _soft_binned_calibration(right, target, 4, boundary_rule) + self.assertAlmostEqual(center_stats[2].sum().item(), 1.0) + self.assertLess(torch.max(torch.abs(left_stats[2] - center_stats[2])).item(), 1e-6) + self.assertLess(torch.max(torch.abs(right_stats[2] - center_stats[2])).item(), 1e-6) + + +class TestCalibrationLoss(unittest.TestCase): + loss_types = (HardL1ACELoss, SoftL1ACELoss) + + def test_manual_loss_values(self): + prediction = torch.tensor([[[0.1, 0.3, 0.7, 0.9]]]) + target = torch.tensor([[[0.0, 0.0, 1.0, 1.0]]]) + result = HardL1ACELoss(num_bins=5)(prediction, target) + torch.testing.assert_close(result, torch.tensor(0.2)) + + prediction = torch.tensor([[[0.0, 0.25, 0.5, 0.75, 1.0]]], dtype=torch.float64) + target = torch.tensor([[[0.0, 0.0, 1.0, 1.0, 1.0]]], dtype=torch.float64) + result = SoftL1ACELoss(num_bins=2)(prediction, target) + torch.testing.assert_close(result, torch.tensor(0.1, dtype=torch.float64)) + + def test_calibrated_and_miscalibrated_values(self): + target = torch.tensor([[[0.0, 1.0, 0.0, 1.0]]]) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + loss = loss_type(num_bins=4) + torch.testing.assert_close(loss(target, target), torch.tensor(0.0)) + self.assertGreater(loss(1.0 - target, target).item(), 0.0) + underconfident = target * 0.5 + 0.25 + self.assertGreater(loss(underconfident, target).item(), 0.0) + + def test_batch_class_reductions_and_weights(self): + prediction = torch.tensor([[[0.1, 0.9], [0.2, 0.8]], [[0.3, 0.7], [0.4, 0.6]]]) + target = torch.tensor([[[0.0, 1.0], [0.0, 1.0]], [[0.0, 1.0], [0.0, 1.0]]]) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + none = loss_type(num_bins=5, reduction="none")(prediction, target) + self.assertEqual(none.shape, (2, 2, 1)) + torch.testing.assert_close(loss_type(num_bins=5)(prediction, target), none.mean()) + torch.testing.assert_close(loss_type(num_bins=5, reduction="sum")(prediction, target), none.sum()) + torch.testing.assert_close(loss_type(num_bins=5, weight=2.0)(prediction, target), 2.0 * none.mean()) + weighted = none * torch.tensor([1.0, 3.0]).view(1, 2, 1) + torch.testing.assert_close( + loss_type(num_bins=5, weight=[1.0, 3.0])(prediction, target), weighted.mean() + ) + + def test_empty_class_is_zero_before_mean(self): + prediction = torch.tensor([[[0.1, 0.9], [0.2, 0.8]]]) + target = torch.tensor([[[0.0, 1.0], [0.0, 0.0]]]) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + none = loss_type(num_bins=5, reduction="none", ignore_empty_classes=True)(prediction, target) + torch.testing.assert_close(none[:, 1], torch.zeros_like(none[:, 1])) + torch.testing.assert_close( + loss_type(num_bins=5, ignore_empty_classes=True)(prediction, target), none.mean() + ) + torch.testing.assert_close(none[:, 0].mean(), 2 * none.mean()) + included = loss_type(num_bins=5, ignore_empty_classes=False)(prediction, target) + self.assertGreater(included.item(), 0.0) + + def test_every_class_empty_returns_differentiable_zero(self): + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + prediction = torch.full((2, 2, 4), 0.25, requires_grad=True) + target = torch.zeros_like(prediction) + result = loss_type(ignore_empty_classes=True)(prediction, target) + torch.testing.assert_close(result, torch.tensor(0.0)) + result.backward() + self.assertIsNotNone(prediction.grad) + torch.testing.assert_close(prediction.grad, torch.zeros_like(prediction.grad)) + + def test_background_exclusion_and_shape(self): + prediction = torch.tensor([[[0.9, 0.1], [0.1, 0.9]]]) + target = torch.tensor([[[1.0, 0.0], [0.0, 1.0]]]) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + result = loss_type(include_background=False, reduction="none")(prediction, target) + self.assertEqual(result.shape, (1, 1, 1)) + + def test_activation_and_one_hot(self): + logits = torch.tensor([[[0.2, -0.4], [-0.2, 0.4]]]) + labels = torch.tensor([[[0, 1]]]) + one_hot_target = torch.nn.functional.one_hot(labels[:, 0], num_classes=2).movedim(-1, 1).float() + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + expected_softmax = loss_type()(torch.softmax(logits, 1), one_hot_target) + torch.testing.assert_close(loss_type(softmax=True, to_onehot_y=True)(logits, labels), expected_softmax) + expected_sigmoid = loss_type()(torch.sigmoid(logits[:, :1]), one_hot_target[:, :1]) + torch.testing.assert_close( + loss_type(sigmoid=True)(logits[:, :1], one_hot_target[:, :1]), expected_sigmoid + ) + torch.testing.assert_close( + loss_type(other_act=lambda value: value.square())(logits.abs(), one_hot_target), + loss_type()(logits.square(), one_hot_target), + ) + + def test_invalid_options_and_shapes(self): + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + with self.assertRaises(ValueError): + loss_type(num_bins=0) + with self.assertRaises(ValueError): + loss_type(sigmoid=True, softmax=True) + with self.assertRaises(TypeError): + loss_type(other_act=1) # type: ignore[arg-type] + with self.assertRaises(ValueError): + loss_type(weight=[1.0, -1.0]) + with self.assertRaises(ValueError): + loss_type(weight=[[1.0]]) + with self.assertRaises(ValueError): + loss_type(reduction="unsupported") + with self.assertRaises(ValueError): + loss_type(weight=[1.0, 2.0, 3.0])(torch.ones(1, 2, 2), torch.ones(1, 2, 2)) + with self.assertRaises(AssertionError): + loss_type()(torch.ones(1, 2, 2), torch.ones(1, 1, 2)) + with self.assertRaises(ValueError): + loss_type()(torch.ones(1, 2), torch.ones(1, 2)) + with self.assertRaises(TypeError): + loss_type()(torch.ones(1, 2, 2, dtype=torch.int64), torch.ones(1, 2, 2)) + with self.assertRaises(ValueError): + SoftL1ACELoss(empty_weight=-1) + + def test_single_channel_warnings(self): + prediction = torch.tensor([[[0.2, 0.8]]]) + target = torch.tensor([[[0.0, 1.0]]]) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + loss_type(softmax=True, to_onehot_y=True, include_background=False)(prediction, target) + self.assertEqual(len(caught), 3) + + def test_soft_empty_weight_masks_bins(self): + prediction = torch.tensor([[[0.25, 0.75]]], requires_grad=True) + target = torch.tensor([[[0.0, 1.0]]]) + result = SoftL1ACELoss(num_bins=2, empty_weight=1.1, ignore_empty_classes=False)(prediction, target) + torch.testing.assert_close(result, torch.tensor(0.0)) + result.backward() + torch.testing.assert_close(prediction.grad, torch.zeros_like(prediction.grad)) + + def test_soft_fractional_counts_are_normalized(self): + prediction = torch.tensor([[[0.5]]]) + target = torch.tensor([[[1.0]]]) + result = SoftL1ACELoss(num_bins=2, empty_weight=0.0)(prediction, target) + torch.testing.assert_close(result, torch.tensor(0.5)) + + def test_dtype_and_gradcheck(self): + target = torch.tensor([[[0.0, 1.0, 0.0, 1.0]]], dtype=torch.float64) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + prediction = torch.tensor([[[0.13, 0.31, 0.67, 0.89]]], dtype=torch.float64, requires_grad=True) + result = loss_type(num_bins=5, ignore_empty_classes=False)(prediction, target) + self.assertEqual(result.dtype, torch.float64) + self.assertTrue( + torch.autograd.gradcheck(loss_type(num_bins=5, ignore_empty_classes=False), (prediction, target)) + ) + result.backward() + self.assertTrue(torch.isfinite(prediction.grad).all()) + self.assertTrue(torch.any(prediction.grad != 0)) + + def test_lower_precision_accumulates_in_float32(self): + prediction = torch.tensor([[[0.1, 0.9]]], dtype=torch.float16) + target = torch.tensor([[[0.0, 1.0]]], dtype=torch.float16) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + result = loss_type(ignore_empty_classes=False)(prediction, target) + self.assertEqual(result.dtype, torch.float32) + self.assertTrue(torch.isfinite(result)) + + def test_3d_spatial_input(self): + prediction = torch.rand(2, 3, 4, 5, 6) + labels = torch.randint(0, 3, (2, 1, 4, 5, 6)) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + result = loss_type(to_onehot_y=True, softmax=True)(prediction, labels) + self.assertEqual(result.shape, torch.Size([])) + self.assertTrue(torch.isfinite(result)) + + def test_script_save(self): + prediction = torch.rand(2, 3, 4, 5, 6) + target = torch.nn.functional.one_hot(torch.randint(0, 3, (2, 4, 5, 6)), num_classes=3) + target = target.movedim(-1, 1).float() + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + check_script_save(loss_type(num_bins=5), prediction, target) + + @skip_if_no_cuda + def test_cuda_float32_and_float64(self): + for dtype in (torch.float32, torch.float64): + for loss_type in self.loss_types: + with self.subTest(dtype=dtype, loss=loss_type.__name__): + prediction = torch.tensor([[[0.1, 0.9]]], device="cuda", dtype=dtype, requires_grad=True) + target = torch.tensor([[[0.0, 1.0]]], device="cuda", dtype=dtype) + result = loss_type(ignore_empty_classes=False).cuda()(prediction, target) + self.assertEqual(result.device.type, "cuda") + self.assertEqual(result.dtype, dtype) + result.backward() + self.assertTrue(torch.isfinite(prediction.grad).all()) + + def test_input_is_not_mutated(self): + prediction = torch.tensor([[[0.1, 0.9]]]) + target = torch.tensor([[[0.0, 1.0]]]) + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + prediction_before = prediction.clone() + target_before = target.clone() + weight = torch.tensor([2.0]) + loss = loss_type(weight=weight) + loss(prediction, target) + torch.testing.assert_close(prediction, prediction_before) + torch.testing.assert_close(target, target_before) + torch.testing.assert_close(loss.class_weight, weight) + + def test_target_has_no_gradient(self): + for loss_type in self.loss_types: + with self.subTest(loss=loss_type.__name__): + prediction = torch.tensor([[[0.1, 0.9]]], requires_grad=True) + target = torch.tensor([[[0.0, 1.0]]], requires_grad=True) + loss_type(ignore_empty_classes=False)(prediction, target).backward() + self.assertIsNotNone(prediction.grad) + self.assertIsNone(target.grad) + + +if __name__ == "__main__": + unittest.main() From 8526ca771d76039f4108acb67e74f405e5acd77e Mon Sep 17 00:00:00 2001 From: Theo Barfoot Date: Thu, 3 Sep 2026 15:50:28 +0000 Subject: [PATCH 2/2] Validate calibration loss weights Signed-off-by: Theo Barfoot --- monai/losses/calibration.py | 6 ++++-- tests/losses/test_calibration_loss.py | 13 +++++++++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/monai/losses/calibration.py b/monai/losses/calibration.py index 5e8b5b0751a..1785b46b5f1 100644 --- a/monai/losses/calibration.py +++ b/monai/losses/calibration.py @@ -109,6 +109,8 @@ def __init__( if class_weight is not None: if class_weight.ndim > 1: raise ValueError("weight must be a scalar or a one-dimensional sequence.") + if not torch.all(torch.isfinite(class_weight)): + raise ValueError("weight must contain only finite values.") if torch.any(class_weight < 0): raise ValueError("the value/values of the `weight` should be no less than 0.") @@ -353,8 +355,8 @@ def __init__( right: bool = False, ignore_empty_classes: bool = True, ) -> None: - if empty_weight < 0: - raise ValueError(f"empty_weight must be >= 0, got {empty_weight}.") + if not 0 <= empty_weight < float("inf"): + raise ValueError(f"empty_weight must be finite and >= 0, got {empty_weight}.") super().__init__( num_bins, include_background, diff --git a/tests/losses/test_calibration_loss.py b/tests/losses/test_calibration_loss.py index d71775268f4..e66818154d5 100644 --- a/tests/losses/test_calibration_loss.py +++ b/tests/losses/test_calibration_loss.py @@ -159,6 +159,14 @@ def test_invalid_options_and_shapes(self): loss_type(other_act=1) # type: ignore[arg-type] with self.assertRaises(ValueError): loss_type(weight=[1.0, -1.0]) + for invalid_weight in (float("nan"), float("inf"), -float("inf")): + with self.subTest(invalid_weight=invalid_weight): + with self.assertRaises(ValueError): + loss_type(weight=invalid_weight) + with self.assertRaises(ValueError): + loss_type(weight=[1.0, invalid_weight]) + with self.assertRaises(ValueError): + loss_type(weight=torch.tensor([invalid_weight])) with self.assertRaises(ValueError): loss_type(weight=[[1.0]]) with self.assertRaises(ValueError): @@ -171,8 +179,9 @@ def test_invalid_options_and_shapes(self): loss_type()(torch.ones(1, 2), torch.ones(1, 2)) with self.assertRaises(TypeError): loss_type()(torch.ones(1, 2, 2, dtype=torch.int64), torch.ones(1, 2, 2)) - with self.assertRaises(ValueError): - SoftL1ACELoss(empty_weight=-1) + for invalid_empty_weight in (-1.0, float("nan"), float("inf"), -float("inf")): + with self.subTest(invalid_empty_weight=invalid_empty_weight), self.assertRaises(ValueError): + SoftL1ACELoss(empty_weight=invalid_empty_weight) def test_single_channel_warnings(self): prediction = torch.tensor([[[0.2, 0.8]]])