Adjust execution order of activation and masking in MaskedDiceLoss

[ytl0623]

Fixes #8655

Description

1. Modified MaskedDiceLoss.forward to ensure masked regions result in 0.0 probability instead of 0.5.
2. Updated unit tests with new expected values, as the fix mathematically changes the loss result by removing the erroneous background contribution.

Types of changes

• 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).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

📝 Walkthrough

Walkthrough

MaskedDiceLoss now exposes an explicit constructor with activation options (sigmoid, softmax, other_act) instead of *args/**kwargs. It validates activation compatibility and that other_act is callable, stores pre-activation flags (pre_sigmoid, pre_softmax, pre_other_act), and initializes the base DiceLoss with activations disabled. The forward path applies the chosen pre-activation (sigmoid, conditional softmax with a single-channel warning, or other_act) before masking and then delegates to the masked spatial Dice computation. Unit tests for MaskedDiceLoss were updated with revised expected numeric outputs.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Linked Issues check	❓ Inconclusive	PR addresses issue #8655 by applying activation before masking; however, implementation deviates from reviewer's recommended approach of storing pre-activation flags and using explicit constructor parameters.	Refactor to store pre-activation flags (pre_sigmoid, pre_softmax, pre_other_act) and use explicit constructor parameters instead of *args/**kwargs as suggested in reviewer feedback.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title accurately describes the main change: adjusting the execution order of activation and masking in MaskedDiceLoss.
Description check	✅ Passed	Description covers the fix, explains test updates, and follows the template with appropriate checklist items marked.
Out of Scope Changes check	✅ Passed	Changes are scoped to MaskedDiceLoss implementation and corresponding test updates; no unrelated modifications detected.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@monai/losses/dice.py`:
- Around line 256-283: The forward block mutates instance flags (self.sigmoid,
self.softmax, self.other_act) before calling spatial_weighted, which is not
thread-safe; instead avoid changing instance state by computing local variables
for the activated input and any activation flags (e.g., use
local_sigmoid/local_softmax/local_other_act or simply pass the already-activated
input) and call spatial_weighted with those locals (or with the activated input
and any needed override parameters) rather than flipping
self.sigmoid/self.softmax/self.other_act; remove the try/finally restore logic
and ensure spatial_weighted signature or call site (spatial_weighted) accepts
the local overrides or only the prepared input to eliminate instance mutation
and race conditions.

🧹 Nitpick comments (3)

monai/losses/dice.py (2)

263-263: Add stacklevel=2 to warning.

Per static analysis (B028): Set stacklevel=2 so the warning points to the caller rather than this line.

Fix

-                warnings.warn("single channel prediction, `softmax=True` ignored.")
+                warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2)

---

257-268: Code duplication with DiceLoss.forward activation handling.

Lines 257-268 mirror DiceLoss.forward (lines 154-165). Consider extracting a _apply_activations helper to reduce duplication.

tests/losses/test_masked_dice_loss.py (1)

115-115: Class name mismatch.

Class TestDiceLoss tests MaskedDiceLoss. Consider renaming to TestMaskedDiceLoss.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

📥 Commits

Reviewing files that changed from the base of the PR and between 57fdd59 and 092990d.

📒 Files selected for processing (2)

• monai/losses/dice.py
• tests/losses/test_masked_dice_loss.py

🧰 Additional context used

📓 Path-based instructions (1)

**/*.py

⚙️ CodeRabbit configuration file

Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. 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. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.

Files:

• tests/losses/test_masked_dice_loss.py
• monai/losses/dice.py

🪛 Ruff (0.14.11)

monai/losses/dice.py

263-263: No explicit stacklevel keyword argument found

Set stacklevel=2

(B028)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (19)

• GitHub Check: min-dep-pytorch (2.6.0)
• GitHub Check: min-dep-py3 (3.9)
• GitHub Check: min-dep-py3 (3.10)
• GitHub Check: min-dep-pytorch (2.8.0)
• GitHub Check: min-dep-pytorch (2.5.1)
• GitHub Check: min-dep-py3 (3.12)
• GitHub Check: min-dep-pytorch (2.7.1)
• GitHub Check: min-dep-py3 (3.11)
• GitHub Check: min-dep-os (macOS-latest)
• GitHub Check: min-dep-os (windows-latest)
• GitHub Check: min-dep-os (ubuntu-latest)
• GitHub Check: quick-py3 (macOS-latest)
• GitHub Check: quick-py3 (ubuntu-latest)
• GitHub Check: flake8-py3 (mypy)
• GitHub Check: flake8-py3 (pytype)
• GitHub Check: quick-py3 (windows-latest)
• GitHub Check: flake8-py3 (codeformat)
• GitHub Check: build-docs
• GitHub Check: packaging

🔇 Additional comments (1)

tests/losses/test_masked_dice_loss.py (1)

30-30: Updated expected values reflect activation-before-masking fix.

The test values (0.333333, 0.301128, 0.579184) correctly correspond to the implementation change where sigmoid is applied before masking. The code confirms that activation layers execute at lines 257–268, then masked loss computation follows at line 279.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

[ericspod]

Hi @ytl0623 thanks for this fix, but I don't feel this is the right solution for it. The issue here is that the activation needs to happen before masking however the approach here creates a new DiceLoss object and so the members of the class inherited from DiceLoss aren't actually used. This isn't correct object-oriented practice and not consistent with other areas in MONAI. I would suggest that we need to keep new variables such as self.pre_sigmoid to track the original argument values passed in, but then pass along False to the super().__init__ call that's been removed. I'd also suggest replacing *args,**kwargs in the constructor with the actual needed named arguments. The constructor would be something like this:

def __init__(
        self,
        include_background: bool = True,
        to_onehot_y: bool = False,
        sigmoid: bool = False,
        softmax: bool = False,
        other_act: Callable | None = None,
        squared_pred: bool = False,
        jaccard: bool = False,
        reduction: LossReduction | str = LossReduction.MEAN,
        smooth_nr: float = 1e-5,
        smooth_dr: float = 1e-5,
        batch: bool = False,
        weight: Sequence[float] | float | int | torch.Tensor | None = None,
        soft_label: bool = False,
    ) -> None:
        """
        Args follow :py:class:`monai.losses.DiceLoss`.
        """
        self.pre_sigmoid=sigmoid
        self.pre_softmax=softmax
        self.pre_other_act=other_act
        super().__init__(
            include_background=include_background,
            to_onehot_y=to_onehot_y,
            sigmoid=False,
            softmax=False,
            other_act=None,
            squared_pred=squared_pred,
            jaccard=jaccard,
            reduction=reduction,
            smooth_nr=smooth_nr,
            smooth_dr=smooth_dr,
            batch=batch,
            weight=class_weight,
            soft_label=soft_label
        )
        self.spatial_weighted = MaskedLoss(loss=super().forward)

Your forward method would be broadly the same though you'd change the names of the members you'd use to the pre_ ones. This allows the inherited DiceLoss members to be used and will be compatible with future inheritance of this class if it happens. We're changing the constructor arguments to be what they should be for DiceLoss, I don't consider this a deprecating change so no need to use the deprecated decorator.

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@monai/losses/dice.py`:
- Around line 262-269: The constructor currently assigns pre_other_act without
validating other_act, so a non-callable other_act causes a runtime error later;
add a check in the DiceLoss __init__ (before assigning self.pre_other_act) that
if other_act is not None and not callable(other_act) then raise TypeError, and
ensure this validation sits alongside the existing incompatible-flag checks for
sigmoid/softmax (referencing other_act, pre_other_act, pre_sigmoid,
pre_softmax).

🧹 Nitpick comments (1)

monai/losses/dice.py (1)

303-303: Add stacklevel=2 to warnings.warn.

Without it, the warning points to this internal line rather than the caller's code.

Proposed fix

-                warnings.warn("single channel prediction, `softmax=True` ignored.")
+                warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2)