Feat: add ignore_index support to DiceLoss#8969
Conversation
📝 WalkthroughWalkthroughDiceLoss gains an Changes
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant DiceLoss
participant Reduction
Caller->>DiceLoss: forward(input, target)
alt ignore_index set
DiceLoss->>DiceLoss: build valid_mask
DiceLoss->>DiceLoss: mask input and target
end
DiceLoss->>DiceLoss: compute per-class loss
alt class_weight set
DiceLoss->>DiceLoss: apply class_weight to per-class loss
end
DiceLoss->>Reduction: reduce(loss)
Reduction-->>Caller: return loss
Related Issues: Not specified. Related PRs: Not specified. Suggested labels: enhancement, losses Suggested reviewers: Not specified. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
monai/losses/dice.py (2)
98-102: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument raised exceptions in
forward.
forwardraisesAssertionError(Line 143) andValueError(Line 186) but the trimmed docstring only hasArgs. Add aRaises:section.As per path instructions, docstrings should describe each raised exception in the appropriate Google-style section.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/losses/dice.py` around lines 98 - 102, The forward docstring in DiceLoss currently documents only Args, but the method can raise AssertionError and ValueError. Update the docstring in DiceLoss.forward to add a Google-style Raises section that names both exceptions and briefly states the conditions that trigger them, keeping the docs aligned with the existing forward validation logic.Source: Path instructions
72-75: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConstructor docstring isn't Google-style.
Only the new
ignore_indexis documented; the rest defers to "standard MONAI DiceLoss". Per Google style each arg should be listed underArgs:.As per path instructions, docstrings should describe each variable in the appropriate Google-style section.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/losses/dice.py` around lines 72 - 75, The constructor docstring for DiceLoss is not in Google style and only mentions ignore_index, so update the docstring in the DiceLoss initializer to use an Args: section that explicitly documents every parameter instead of referring to “standard MONAI DiceLoss.” Make sure each argument used by the DiceLoss constructor is listed with a brief description, and include ignore_index in the same section alongside the existing parameters.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@monai/losses/dice.py`:
- Around line 127-148: The ignore_index handling in the Dice loss is masking an
entire class channel in the one-hot path instead of only the ignored pixels, so
the mask logic in the Dice class should be changed to produce a spatial valid
mask from the ignored channel rather than zeroing the full channel. Update the
masking in the Dice loss flow so `valid_mask` reflects ignored voxels/pixels per
location and then apply it consistently to both `input` and `target`, keeping
one-hot and index targets equivalent while preserving the existing
`include_background` behavior.
- Around line 116-148: Add unit tests for DiceLoss ignore_index handling in both
index-format B1HW targets and one-hot BNHW targets, including the
include_background=False path and out-of-range ignore_index values. Extend the
existing DiceLoss test coverage to verify ignored regions do not contribute to
the computed loss or backpropagated gradients, and reference the DiceLoss
forward/masking logic when asserting the expected behavior.
---
Nitpick comments:
In `@monai/losses/dice.py`:
- Around line 98-102: The forward docstring in DiceLoss currently documents only
Args, but the method can raise AssertionError and ValueError. Update the
docstring in DiceLoss.forward to add a Google-style Raises section that names
both exceptions and briefly states the conditions that trigger them, keeping the
docs aligned with the existing forward validation logic.
- Around line 72-75: The constructor docstring for DiceLoss is not in Google
style and only mentions ignore_index, so update the docstring in the DiceLoss
initializer to use an Args: section that explicitly documents every parameter
instead of referring to “standard MONAI DiceLoss.” Make sure each argument used
by the DiceLoss constructor is listed with a brief description, and include
ignore_index in the same section alongside the existing parameters.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ed022536-cc87-49a6-861c-d8f8312c39e6
📒 Files selected for processing (1)
monai/losses/dice.py
Signed-off-by: qwepablo12 <grazava6@gmail.com>
e690fbd to
7fdeae6
Compare
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
monai/losses/dice.py (1)
118-131: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSanitize
ignore_indexbeforeone_hotvalid_maskis built too late here; out-of-range ignored labels like255or-100will still hitone_hot(...)and fail inscatter_. Neutralize masked labels before conversion.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/losses/dice.py` around lines 118 - 131, `DiceLoss` builds `valid_mask` in the `ignore_index` path, but `one_hot(...)` still receives the raw target and can fail on ignored labels like 255 or -100. Update the `DiceLoss` target-preprocessing flow so masked labels are neutralized before calling `one_hot`, using the existing `ignore_index`, `valid_mask`, and `to_onehot_y` logic in `dice.py`. Keep the mask behavior intact while ensuring the labels passed into `one_hot` are always in-range.
🧹 Nitpick comments (1)
monai/losses/dice.py (1)
98-102: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring lost its
Returns/Raisessections.forwardraisesAssertionError(Line 142) andValueError(Lines 170, 175, 186) and returns a tensor; document them in Google style.As per path instructions: "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/losses/dice.py` around lines 98 - 102, The Dice loss forward docstring is missing the Google-style Returns and Raises sections even though `forward` returns a tensor and can throw `AssertionError` and `ValueError`. Update the docstring for `forward` in `dice.py` to document the returned tensor and list the exceptions raised by the input validation paths, keeping the existing argument descriptions and using the `forward` method as the anchor.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/losses/test_dice_loss.py`:
- Around line 244-247: The test `test_ignore_index_one_hot_target` uses
`one_hot` but never imports it, so the test will fail with a NameError. Add the
appropriate MONAI import for `one_hot` at the top of the test module, using the
existing test file’s import section and the `one_hot` helper referenced in
`test_ignore_index_one_hot_target`.
---
Outside diff comments:
In `@monai/losses/dice.py`:
- Around line 118-131: `DiceLoss` builds `valid_mask` in the `ignore_index`
path, but `one_hot(...)` still receives the raw target and can fail on ignored
labels like 255 or -100. Update the `DiceLoss` target-preprocessing flow so
masked labels are neutralized before calling `one_hot`, using the existing
`ignore_index`, `valid_mask`, and `to_onehot_y` logic in `dice.py`. Keep the
mask behavior intact while ensuring the labels passed into `one_hot` are always
in-range.
---
Nitpick comments:
In `@monai/losses/dice.py`:
- Around line 98-102: The Dice loss forward docstring is missing the
Google-style Returns and Raises sections even though `forward` returns a tensor
and can throw `AssertionError` and `ValueError`. Update the docstring for
`forward` in `dice.py` to document the returned tensor and list the exceptions
raised by the input validation paths, keeping the existing argument descriptions
and using the `forward` method as the anchor.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 0162add8-4953-424e-bb94-a4068c72c6b4
📒 Files selected for processing (2)
monai/losses/dice.pytests/losses/test_dice_loss.py
| def test_ignore_index_one_hot_target(self): | ||
| input_tensor = torch.randn(2, 4, 8, 8) | ||
| target_idx = torch.randint(0, 4, (2, 1, 8, 8)) | ||
| target_onehot = one_hot(target_idx, num_classes=4) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Missing one_hot import — test will raise NameError.
one_hot is used at Line 247 but never imported. It's a MONAI utility (monai.networks.one_hot / monai.networks.utils.one_hot), not a builtin.
🐛 Proposed fix
+from monai.networks import one_hot🧰 Tools
🪛 Ruff (0.15.20)
[error] 247-247: Undefined name one_hot
(F821)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/losses/test_dice_loss.py` around lines 244 - 247, The test
`test_ignore_index_one_hot_target` uses `one_hot` but never imports it, so the
test will fail with a NameError. Add the appropriate MONAI import for `one_hot`
at the top of the test module, using the existing test file’s import section and
the `one_hot` helper referenced in `test_ignore_index_one_hot_target`.
Source: Linters/SAST tools
Description
Addresses #8734.
This PR introduces
ignore_indexsupport forDiceLoss. Whenignore_indexis specified, the target regions matching this index are masked out and completely excluded from the Dice coefficient computation (both numerator and denominator). This ensures that ignored regions/classes do not affect the loss values or gradients during training.Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.