Skip to content

Commit 7af8653

Browse files
Rusheel86ericspod
andauthored
feat(losses/metrics): implement ignore_index support across dice and … (#8757)
## Fixes #8734 ### Description This PR introduces comprehensive native support for `ignore_index` in core MONAI losses and metrics, as requested in issue #8734. The `ignore_index` parameter allows specific label values (e.g., padding, unlabeled regions, or boundary artifacts) to be excluded from loss and metric calculations which is a critical feature for medical imaging workflows. ### Implementation Summary **Losses with `ignore_index` support:** - `DiceLoss` - masks ignored voxels before dice computation - `FocalLoss` - applies masking after one-hot conversion - `TverskyLoss` - consistent masking approach with DiceLoss - `UnifiedFocalLoss` - handles ignore_index in binary target conversion **Metrics with `ignore_index` support:** - `MeanDice` & `GeneralizedDiceScore` - spatial masking before computation - `MeanIoU` - excludes ignored classes from IoU calculation - `ConfusionMatrixMetric` - filters ignored indices from confusion matrix - `HausdorffDistanceMetric` - masks boundary computation - `SurfaceDiceMetric` & `SurfaceDistanceMetric` - handles ignore_index in surface calculations **Key Fixes:** - Added support for `ignore_index` - Added comprehensive test coverage in `test_ignore_index_losses.py` and `test_ignore_index_metrics.py` - Replaced hardcoded numeric stability constants with configurable `epsilon` parameter in `GeneralizedDiceScore` ### Types of changes - [x] Non-breaking change (fix or new feature that would not break existing functionality) - [ ] Breaking change (fix or new feature that would cause existing functionality to change) - [x] New tests added to cover the changes (`test_ignore_index_losses.py`, `test_ignore_index_metrics.py`) - [x] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage` - [x] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests` - [x] In-line docstrings updated - [x] Documentation updated, tested `make html` command in the `docs/` folder ### Related Issues Could merge with or addresses similar requirements as #8667 (Ignore Class) --------- Signed-off-by: Rusheel Sharma <rusheelhere@gmail.com> Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>
1 parent 453688a commit 7af8653

15 files changed

Lines changed: 959 additions & 106 deletions

monai/losses/dice.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@
2222

2323
from monai.losses.focal_loss import FocalLoss
2424
from monai.losses.spatial_mask import MaskedLoss
25-
from monai.losses.utils import compute_tp_fp_fn
25+
from monai.losses.utils import compute_tp_fp_fn, mask_loss_inputs
26+
from monai.metrics.utils import create_ignore_mask
2627
from monai.networks import one_hot
2728
from monai.utils import DiceCEReduction, LossReduction, Weight, look_up_option
2829

@@ -67,6 +68,7 @@ def __init__(
6768
batch: bool = False,
6869
weight: Sequence[float] | float | int | torch.Tensor | None = None,
6970
soft_label: bool = False,
71+
ignore_index: int | None = None,
7072
) -> None:
7173
"""
7274
Args:
@@ -100,6 +102,10 @@ def __init__(
100102
The value/values should be no less than 0. Defaults to None.
101103
soft_label: whether the target contains non-binary values (soft labels) or not.
102104
If True a soft label formulation of the loss will be used.
105+
ignore_index: single integer class index (or sentinel value) to ignore from the loss computation.
106+
Voxels with this label are excluded from the loss, which is useful for padding, unlabeled regions,
107+
or boundary artifacts. For federated or aggregated settings, ensure all clients use the same
108+
ignore_index to keep loss values comparable.
103109
104110
Raises:
105111
TypeError: When ``other_act`` is not an ``Optional[Callable]``.
@@ -122,6 +128,7 @@ def __init__(
122128
self.smooth_nr = float(smooth_nr)
123129
self.smooth_dr = float(smooth_dr)
124130
self.batch = batch
131+
self.ignore_index = ignore_index
125132
weight = torch.as_tensor(weight) if weight is not None else None
126133
self.register_buffer("class_weight", weight)
127134
self.class_weight: None | torch.Tensor
@@ -163,10 +170,15 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
163170
if self.other_act is not None:
164171
input = self.other_act(input)
165172

173+
original_target = target
174+
166175
if self.to_onehot_y:
167176
if n_pred_ch == 1:
168177
warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2)
169178
else:
179+
if self.ignore_index is not None:
180+
if self.ignore_index < 0 or self.ignore_index >= n_pred_ch:
181+
target = torch.where(target == self.ignore_index, torch.zeros_like(target), target)
170182
target = one_hot(target, num_classes=n_pred_ch)
171183

172184
if not self.include_background:
@@ -177,9 +189,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
177189
target = target[:, 1:]
178190
input = input[:, 1:]
179191

192+
mask = create_ignore_mask(original_target, self.ignore_index)
193+
180194
if target.shape != input.shape:
181195
raise AssertionError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
182196

197+
input, target = mask_loss_inputs(input, target, self.ignore_index, mask=mask)
198+
183199
# reducing only spatial dimensions (not batch nor channels)
184200
reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist()
185201
if self.batch:

monai/losses/focal_loss.py

Lines changed: 58 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import torch.nn.functional as F
1919
from torch.nn.modules.loss import _Loss
2020

21+
from monai.losses.utils import mask_loss_inputs
22+
from monai.metrics.utils import create_ignore_mask
2123
from monai.networks import one_hot
2224
from monai.utils import LossReduction
2325

@@ -73,6 +75,7 @@ def __init__(
7375
weight: Sequence[float] | float | int | torch.Tensor | None = None,
7476
reduction: LossReduction | str = LossReduction.MEAN,
7577
use_softmax: bool = False,
78+
ignore_index: int | None = None,
7679
) -> None:
7780
"""
7881
Args:
@@ -99,6 +102,10 @@ def __init__(
99102
100103
use_softmax: whether to use softmax to transform the original logits into probabilities.
101104
If True, softmax is used. If False, sigmoid is used. Defaults to False.
105+
ignore_index: single integer class index (or sentinel value) to ignore from the loss computation.
106+
Voxels with this label are excluded from the loss, which is useful for padding, unlabeled regions,
107+
or boundary artifacts. For federated or aggregated settings, ensure all clients use the same
108+
ignore_index to keep loss values comparable.
102109
103110
Example:
104111
>>> import torch
@@ -124,6 +131,7 @@ def __init__(
124131
weight = torch.as_tensor(weight) if weight is not None else None
125132
self.register_buffer("class_weight", weight)
126133
self.class_weight: None | torch.Tensor
134+
self.ignore_index = ignore_index
127135

128136
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
129137
"""
@@ -144,10 +152,15 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
144152
"""
145153
n_pred_ch = input.shape[1]
146154

155+
original_target = target
156+
147157
if self.to_onehot_y:
148158
if n_pred_ch == 1:
149159
warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2)
150160
else:
161+
if self.ignore_index is not None:
162+
if self.ignore_index < 0 or self.ignore_index >= n_pred_ch:
163+
target = torch.where(target == self.ignore_index, torch.zeros_like(target), target)
151164
target = one_hot(target, num_classes=n_pred_ch)
152165

153166
if not self.include_background:
@@ -158,9 +171,15 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
158171
target = target[:, 1:]
159172
input = input[:, 1:]
160173

174+
mask = None
175+
if self.ignore_index is not None:
176+
mask = create_ignore_mask(original_target, self.ignore_index)
177+
161178
if target.shape != input.shape:
162179
raise ValueError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
163180

181+
input, target = mask_loss_inputs(input, target, self.ignore_index, mask=mask)
182+
164183
loss: torch.Tensor | None = None
165184
input = input.float()
166185
target = target.float()
@@ -176,24 +195,32 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
176195
else:
177196
loss = sigmoid_focal_loss(input, target, self.gamma, alpha_arg)
178197

179-
num_of_classes = target.shape[1]
180-
if self.class_weight is not None and num_of_classes != 1:
181-
# make sure the lengths of weights are equal to the number of classes
182-
if self.class_weight.ndim == 0:
183-
self.class_weight = torch.as_tensor([self.class_weight] * num_of_classes)
184-
else:
185-
if self.class_weight.shape[0] != num_of_classes:
198+
if mask is not None:
199+
loss = loss * mask
200+
201+
if self.class_weight is not None:
202+
cw = torch.as_tensor(self.class_weight, device=loss.device, dtype=loss.dtype)
203+
num_classes = loss.shape[1]
204+
205+
if cw.ndim > 0:
206+
if num_classes == 1:
207+
raise ValueError("Per-class class_weight is not supported for single-channel outputs.")
208+
if cw.numel() != num_classes:
186209
raise ValueError(
187-
"The length of the `weight` sequence should be the same as the number of classes. "
188-
"If `include_background=False`, the weight should not include the background category class 0."
210+
f"The number of class_weight ({cw.numel()}) must match the number of "
211+
f"output channels ({num_classes})."
189212
)
190-
if self.class_weight.min() < 0:
191-
raise ValueError("the value/values of the `weight` should be no less than 0.")
192-
# apply class_weight to loss
193-
self.class_weight = self.class_weight.to(loss)
194-
broadcast_dims = [-1] + [1] * len(target.shape[2:])
195-
self.class_weight = self.class_weight.view(broadcast_dims)
196-
loss = self.class_weight * loss
213+
if (cw < 0).any():
214+
raise ValueError("class_weight values must be non-negative.")
215+
else:
216+
if cw < 0:
217+
raise ValueError("class_weight values must be non-negative.")
218+
219+
if cw.ndim == 0:
220+
loss = loss * cw
221+
else:
222+
broadcast_shape = [1, num_classes] + [1] * (loss.ndim - 2)
223+
loss = loss * cw.view(broadcast_shape)
197224

198225
if self.reduction == LossReduction.SUM.value:
199226
# 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:
202229
# parameterize if necessary. (Or justify why the mean should be there)
203230
average_spatial_dims = True
204231
if average_spatial_dims:
205-
loss = loss.mean(dim=list(range(2, len(target.shape))))
232+
if mask is not None:
233+
spatial_dims = list(range(2, len(target.shape)))
234+
sum_mask = mask.sum(dim=spatial_dims, keepdim=True)
235+
loss = loss.sum(dim=spatial_dims, keepdim=True) / sum_mask.clamp(min=torch.finfo(mask.dtype).eps)
236+
else:
237+
loss = loss.mean(dim=list(range(2, len(target.shape))))
206238
loss = loss.sum()
239+
207240
elif self.reduction == LossReduction.MEAN.value:
208-
loss = loss.mean()
241+
if mask is not None:
242+
# Sum loss over valid (non-ignored) elements, then divide by mask count
243+
loss = loss.sum() / mask.sum().clamp(min=torch.finfo(mask.dtype).eps)
244+
else:
245+
loss = loss.mean()
246+
209247
elif self.reduction == LossReduction.NONE.value:
210248
pass
249+
211250
else:
212251
raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].')
252+
213253
return loss
214254

215255

monai/losses/tversky.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@
1717
import torch
1818
from torch.nn.modules.loss import _Loss
1919

20-
from monai.losses.utils import compute_tp_fp_fn
20+
from monai.losses.utils import compute_tp_fp_fn, mask_loss_inputs
21+
from monai.metrics.utils import create_ignore_mask
2122
from monai.networks import one_hot
2223
from monai.utils import LossReduction
2324

@@ -51,6 +52,7 @@ def __init__(
5152
smooth_dr: float = 1e-5,
5253
batch: bool = False,
5354
soft_label: bool = False,
55+
ignore_index: int | None = None,
5456
) -> None:
5557
"""
5658
Args:
@@ -77,6 +79,10 @@ def __init__(
7779
before any `reduction`.
7880
soft_label: whether the target contains non-binary values (soft labels) or not.
7981
If True a soft label formulation of the loss will be used.
82+
ignore_index: single integer class index (or sentinel value) to ignore from the loss computation.
83+
Voxels with this label are excluded from the loss, which is useful for padding, unlabeled regions,
84+
or boundary artifacts. For federated or aggregated settings, ensure all clients use the same
85+
ignore_index to keep loss values comparable.
8086
8187
Raises:
8288
TypeError: When ``other_act`` is not an ``Optional[Callable]``.
@@ -101,6 +107,7 @@ def __init__(
101107
self.smooth_dr = float(smooth_dr)
102108
self.batch = batch
103109
self.soft_label = soft_label
110+
self.ignore_index = ignore_index
104111

105112
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
106113
"""
@@ -125,12 +132,21 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
125132
if self.other_act is not None:
126133
input = self.other_act(input)
127134

135+
original_target = target
136+
128137
if self.to_onehot_y:
129138
if n_pred_ch == 1:
130139
warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2)
131140
else:
141+
if self.ignore_index is not None:
142+
if self.ignore_index < 0 or self.ignore_index >= n_pred_ch:
143+
target = torch.where(target == self.ignore_index, torch.zeros_like(target), target)
132144
target = one_hot(target, num_classes=n_pred_ch)
133145

146+
if self.ignore_index is not None:
147+
mask = create_ignore_mask(original_target, self.ignore_index)
148+
input, target = mask_loss_inputs(input, target, self.ignore_index, mask=mask)
149+
134150
if not self.include_background:
135151
if n_pred_ch == 1:
136152
warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2)

0 commit comments

Comments
 (0)