feat(losses/metrics): implement ignore_index support across dice and …

monai/losses/dice.py

@@ -22,7 +22,8 @@
 
 from monai.losses.focal_loss import FocalLoss
 from monai.losses.spatial_mask import MaskedLoss
-from monai.losses.utils import compute_tp_fp_fn
+from monai.losses.utils import compute_tp_fp_fn, mask_loss_inputs
+from monai.metrics.utils import create_ignore_mask
 from monai.networks import one_hot
 from monai.utils import DiceCEReduction, LossReduction, Weight, look_up_option
 
@@ -67,6 +68,7 @@ def __init__(
         batch: bool = False,
         weight: Sequence[float] | float | int | torch.Tensor | None = None,
         soft_label: bool = False,
+        ignore_index: int | None = None,
     ) -> None:
         """
         Args:
@@ -100,6 +102,10 @@ def __init__(
                 The value/values should be no less than 0. Defaults to None.
             soft_label: whether the target contains non-binary values (soft labels) or not.
                 If True a soft label formulation of the loss will be used.
+            ignore_index: single integer class index (or sentinel value) to ignore from the loss computation.
+                Voxels with this label are excluded from the loss, which is useful for padding, unlabeled regions,
+                or boundary artifacts. For federated or aggregated settings, ensure all clients use the same
+                ignore_index to keep loss values comparable.
 
         Raises:
             TypeError: When ``other_act`` is not an ``Optional[Callable]``.
@@ -122,6 +128,7 @@ def __init__(
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
         self.batch = batch
+        self.ignore_index = ignore_index
         weight = torch.as_tensor(weight) if weight is not None else None
         self.register_buffer("class_weight", weight)
         self.class_weight: None | torch.Tensor
@@ -163,10 +170,15 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         if self.other_act is not None:
             input = self.other_act(input)
 
+        original_target = target
+
         if self.to_onehot_y:
             if n_pred_ch == 1:
                 warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2)
             else:
+                if self.ignore_index is not None:
+                    if self.ignore_index < 0 or self.ignore_index >= n_pred_ch:
+                        target = torch.where(target == self.ignore_index, torch.zeros_like(target), target)
                 target = one_hot(target, num_classes=n_pred_ch)
 
         if not self.include_background:
@@ -177,9 +189,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                 target = target[:, 1:]
                 input = input[:, 1:]
 
+        mask = create_ignore_mask(original_target, self.ignore_index)
+
         if target.shape != input.shape:
             raise AssertionError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
 
+        input, target = mask_loss_inputs(input, target, self.ignore_index, mask=mask)
+
         # reducing only spatial dimensions (not batch nor channels)
         reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist()
         if self.batch:

monai/losses/focal_loss.py

@@ -18,6 +18,8 @@
 import torch.nn.functional as F
 from torch.nn.modules.loss import _Loss
 
+from monai.losses.utils import mask_loss_inputs
+from monai.metrics.utils import create_ignore_mask
 from monai.networks import one_hot
 from monai.utils import LossReduction
 
@@ -73,6 +75,7 @@ def __init__(
         weight: Sequence[float] | float | int | torch.Tensor | None = None,
         reduction: LossReduction | str = LossReduction.MEAN,
         use_softmax: bool = False,
+        ignore_index: int | None = None,
     ) -> None:
         """
         Args:
@@ -99,6 +102,10 @@ def __init__(
 
             use_softmax: whether to use softmax to transform the original logits into probabilities.
                 If True, softmax is used. If False, sigmoid is used. Defaults to False.
+            ignore_index: single integer class index (or sentinel value) to ignore from the loss computation.
+                Voxels with this label are excluded from the loss, which is useful for padding, unlabeled regions,
+                or boundary artifacts. For federated or aggregated settings, ensure all clients use the same
+                ignore_index to keep loss values comparable.
 
         Example:
             >>> import torch
@@ -124,6 +131,7 @@ def __init__(
         weight = torch.as_tensor(weight) if weight is not None else None
         self.register_buffer("class_weight", weight)
         self.class_weight: None | torch.Tensor
+        self.ignore_index = ignore_index
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -144,10 +152,15 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
         n_pred_ch = input.shape[1]
 
+        original_target = target
+
         if self.to_onehot_y:
             if n_pred_ch == 1:
                 warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2)
             else:
+                if self.ignore_index is not None:
+                    if self.ignore_index < 0 or self.ignore_index >= n_pred_ch:
+                        target = torch.where(target == self.ignore_index, torch.zeros_like(target), target)
                 target = one_hot(target, num_classes=n_pred_ch)
 
         if not self.include_background:
@@ -158,9 +171,15 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                 target = target[:, 1:]
                 input = input[:, 1:]
 
+        mask = None
+        if self.ignore_index is not None:
+            mask = create_ignore_mask(original_target, self.ignore_index)
+
         if target.shape != input.shape:
             raise ValueError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
 
+        input, target = mask_loss_inputs(input, target, self.ignore_index, mask=mask)
+
         loss: torch.Tensor | None = None
         input = input.float()
         target = target.float()
@@ -176,24 +195,32 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         else:
             loss = sigmoid_focal_loss(input, target, self.gamma, alpha_arg)
 
-        num_of_classes = target.shape[1]
-        if self.class_weight is not None and num_of_classes != 1:
-            # make sure the lengths of weights are equal to the number of classes
-            if self.class_weight.ndim == 0:
-                self.class_weight = torch.as_tensor([self.class_weight] * num_of_classes)
-            else:
-                if self.class_weight.shape[0] != num_of_classes:
+        if mask is not None:
+            loss = loss * mask
+
+        if self.class_weight is not None:
+            cw = torch.as_tensor(self.class_weight, device=loss.device, dtype=loss.dtype)
+            num_classes = loss.shape[1]
+
+            if cw.ndim > 0:
+                if num_classes == 1:
+                    raise ValueError("Per-class class_weight is not supported for single-channel outputs.")
+                if cw.numel() != num_classes:
                     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."
+                        f"The number of class_weight ({cw.numel()}) must match the number of "
+                        f"output channels ({num_classes})."
                     )
-            if self.class_weight.min() < 0:
-                raise ValueError("the value/values of the `weight` should be no less than 0.")
-            # apply class_weight to loss
-            self.class_weight = self.class_weight.to(loss)
-            broadcast_dims = [-1] + [1] * len(target.shape[2:])
-            self.class_weight = self.class_weight.view(broadcast_dims)
-            loss = self.class_weight * loss
+                if (cw < 0).any():
+                    raise ValueError("class_weight values must be non-negative.")
+            else:
+                if cw < 0:
+                    raise ValueError("class_weight values must be non-negative.")
+
+            if cw.ndim == 0:
+                loss = loss * cw
+            else:
+                broadcast_shape = [1, num_classes] + [1] * (loss.ndim - 2)
+                loss = loss * cw.view(broadcast_shape)
 
         if self.reduction == LossReduction.SUM.value:
             # Previously there was a mean over the last dimension, which did not
@@ -202,14 +229,27 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             # parameterize if necessary. (Or justify why the mean should be there)
             average_spatial_dims = True
             if average_spatial_dims:
-                loss = loss.mean(dim=list(range(2, len(target.shape))))
+                if mask is not None:
+                    spatial_dims = list(range(2, len(target.shape)))
+                    sum_mask = mask.sum(dim=spatial_dims, keepdim=True)
+                    loss = loss.sum(dim=spatial_dims, keepdim=True) / sum_mask.clamp(min=torch.finfo(mask.dtype).eps)
+                else:
+                    loss = loss.mean(dim=list(range(2, len(target.shape))))
             loss = loss.sum()
+
         elif self.reduction == LossReduction.MEAN.value:
-            loss = loss.mean()
+            if mask is not None:
+                # Sum loss over valid (non-ignored) elements, then divide by mask count
+                loss = loss.sum() / mask.sum().clamp(min=torch.finfo(mask.dtype).eps)
+            else:
+                loss = loss.mean()
+
         elif self.reduction == LossReduction.NONE.value:
             pass
+
         else:
             raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].')
+
         return loss
 
 

monai/losses/tversky.py

@@ -17,7 +17,8 @@
 import torch
 from torch.nn.modules.loss import _Loss
 
-from monai.losses.utils import compute_tp_fp_fn
+from monai.losses.utils import compute_tp_fp_fn, mask_loss_inputs
+from monai.metrics.utils import create_ignore_mask
 from monai.networks import one_hot
 from monai.utils import LossReduction
 
@@ -51,6 +52,7 @@ def __init__(
         smooth_dr: float = 1e-5,
         batch: bool = False,
         soft_label: bool = False,
+        ignore_index: int | None = None,
     ) -> None:
         """
         Args:
@@ -77,6 +79,10 @@ def __init__(
                 before any `reduction`.
             soft_label: whether the target contains non-binary values (soft labels) or not.
                 If True a soft label formulation of the loss will be used.
+            ignore_index: single integer class index (or sentinel value) to ignore from the loss computation.
+                Voxels with this label are excluded from the loss, which is useful for padding, unlabeled regions,
+                or boundary artifacts. For federated or aggregated settings, ensure all clients use the same
+                ignore_index to keep loss values comparable.
 
         Raises:
             TypeError: When ``other_act`` is not an ``Optional[Callable]``.
@@ -101,6 +107,7 @@ def __init__(
         self.smooth_dr = float(smooth_dr)
         self.batch = batch
         self.soft_label = soft_label
+        self.ignore_index = ignore_index
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -125,12 +132,21 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         if self.other_act is not None:
             input = self.other_act(input)
 
+        original_target = target
+
         if self.to_onehot_y:
             if n_pred_ch == 1:
                 warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2)
             else:
+                if self.ignore_index is not None:
+                    if self.ignore_index < 0 or self.ignore_index >= n_pred_ch:
+                        target = torch.where(target == self.ignore_index, torch.zeros_like(target), target)
                 target = one_hot(target, num_classes=n_pred_ch)
 
+        if self.ignore_index is not None:
+            mask = create_ignore_mask(original_target, self.ignore_index)
+            input, target = mask_loss_inputs(input, target, self.ignore_index, mask=mask)
+
         if not self.include_background:
             if n_pred_ch == 1:
                 warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2)