Add sigmoid/softmax activation to AsymmetricUnifiedFocalLoss#8863
Open
AlexanderSanin wants to merge 5 commits into
Open
Add sigmoid/softmax activation to AsymmetricUnifiedFocalLoss#8863AlexanderSanin wants to merge 5 commits into
AlexanderSanin wants to merge 5 commits into
Conversation
) Catching BaseException inadvertently suppresses KeyboardInterrupt, SystemExit, and GeneratorExit, which should nearly always propagate. All 17 occurrences across monai/ and tests/ are replaced with Exception, which is the appropriate base class for catchable errors. Signed-off-by: Oleksandr Sanin <alexaaander.sanin@gmail.com>
…-MONAI#8603) Add use_softmax and use_sigmoid parameters so users can pass raw logits directly. When both are False (default), the input is assumed to already be probabilities, preserving backward compatibility. Also removes the stale TODO comment about multi-class support and adds proper docstrings for the reduction parameter. Signed-off-by: Oleksandr Sanin <alexaaander.sanin@gmail.com>
AlexanderSanin
requested review from
KumoLiu,
Nic-Ma,
ericspod and
virginiafdez
as code owners
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR systematically narrows many catch/suppress blocks from catching BaseException to Exception across multiple modules and tests, and adds two constructor flags (use_softmax, use_sigmoid) to AsymmetricUnifiedFocalLoss with runtime mutual-exclusivity enforcement and accompanying unit tests that validate activations and the exclusivity check. Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
tests/losses/test_unified_focal_loss.py (1)
71-76:⚠️ Potential issue | 🟠 Major | ⚡ Quick winSame issues as
test_use_sigmoid.Single-channel logits will fail one-hot conversion, and the assertion is too weak to verify softmax application.
🤖 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_unified_focal_loss.py` around lines 71 - 76, test_use_softmax fails because single-channel logits break one-hot conversion and the current assertion is too weak; update the test to feed multi-channel logits (C=2) and matching one-hot y_true so AsymmetricUnifiedFocalLoss(use_softmax=True) can perform softmax, then strengthen the assertion by comparing the softmax-enabled loss to a baseline (e.g., compute loss_softmax = loss(y_pred, y_true) with AsymmetricUnifiedFocalLoss(use_softmax=True) and loss_sigmoid = AsymmetricUnifiedFocalLoss(use_softmax=False)(same y_pred, y_true)) and assert both are finite and that loss_softmax differs from loss_sigmoid (or is less than, depending on expected behavior) to verify softmax was applied; reference test_use_softmax, AsymmetricUnifiedFocalLoss, use_softmax, y_pred and y_true when making the changes.
🤖 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/unified_focal_loss.py`:
- Around line 241-244: The activation is applied after converting single-channel
y_pred to one-hot causing integer-index errors for single-channel logits; update
the UnifiedFocalLoss forward (or the method handling y_pred) to either apply
softmax/sigmoid before the one_hot conversion or add an explicit input
validation that raises a clear error when y_pred has a single channel while
use_softmax or use_sigmoid is True (referencing variables y_pred, use_softmax,
use_sigmoid and the one_hot conversion block) so users get a helpful message
rather than a silent failure.
- Around line 241-244: The code enables softmax via use_softmax but the
component losses AsymmetricFocalLoss and AsymmetricFocalTverskyLoss index
hardcoded class channels ([:, 0] and [:, 1]) so softmax with >2 channels will be
broken; update the functions to validate and reject multi-class inputs (e.g.,
check y_pred.shape[1] or num_classes and raise a clear ValueError if != 2) or
explicitly document in the use_softmax docstring that only binary (2-class)
predictions are supported, and ensure the error message references use_softmax,
AsymmetricFocalLoss, and AsymmetricFocalTverskyLoss so callers know why their
multi-channel input is unsupported.
In `@tests/losses/test_unified_focal_loss.py`:
- Around line 64-69: The test crashes because AsymmetricUnifiedFocalLoss's
internal one_hot conversion is called when y_pred is single-channel float logits
([2,1,2,2]) and those logits are being cast to long, producing out-of-range
indices; change the test to provide multi-channel logits (e.g. shape [2,2,2,2])
when use_sigmoid=True or else supply integer class indices for y_true so one_hot
isn't fed raw logits; update the test_use_sigmoid to create y_pred with two
channels and matching y_true (or use class indices) so one_hot/scatter_ receives
valid class indices.
---
Duplicate comments:
In `@tests/losses/test_unified_focal_loss.py`:
- Around line 71-76: test_use_softmax fails because single-channel logits break
one-hot conversion and the current assertion is too weak; update the test to
feed multi-channel logits (C=2) and matching one-hot y_true so
AsymmetricUnifiedFocalLoss(use_softmax=True) can perform softmax, then
strengthen the assertion by comparing the softmax-enabled loss to a baseline
(e.g., compute loss_softmax = loss(y_pred, y_true) with
AsymmetricUnifiedFocalLoss(use_softmax=True) and loss_sigmoid =
AsymmetricUnifiedFocalLoss(use_softmax=False)(same y_pred, y_true)) and assert
both are finite and that loss_softmax differs from loss_sigmoid (or is less
than, depending on expected behavior) to verify softmax was applied; reference
test_use_softmax, AsymmetricUnifiedFocalLoss, use_softmax, y_pred and y_true
when making the changes.
🪄 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: 0e334f54-a695-4d0b-bc39-462797493bc0
📒 Files selected for processing (15)
monai/__init__.pymonai/apps/auto3dseg/data_analyzer.pymonai/apps/auto3dseg/ensemble_builder.pymonai/apps/auto3dseg/utils.pymonai/apps/detection/metrics/coco.pymonai/apps/nnunet/nnunetv2_runner.pymonai/config/deviceconfig.pymonai/data/__init__.pymonai/inferers/inferer.pymonai/losses/unified_focal_loss.pymonai/utils/tf32.pytests/apps/detection/networks/test_retinanet.pytests/losses/test_unified_focal_loss.pytests/networks/nets/test_resnet.pytests/test_utils.py
Comment on lines
+241
to
+244
| if self.use_softmax: | ||
| y_pred = torch.softmax(y_pred, dim=1) | ||
| elif self.use_sigmoid: | ||
| y_pred = torch.sigmoid(y_pred) |
Contributor
There was a problem hiding this comment.
Activation applied after one-hot conversion breaks single-channel logit inputs.
Lines 227-229 convert single-channel y_pred to one-hot before activation is applied here. This means:
- Single-channel inputs are treated as class indices (discrete values), not logits
- Passing single-channel logits with
use_sigmoid=Truewill fail becauseone_hot()expects integers - Users must pass 2-channel inputs to use the activation flags
Document this requirement or add validation to raise a clear error for single-channel inputs when activation flags are set.
🤖 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/unified_focal_loss.py` around lines 241 - 244, The activation is
applied after converting single-channel y_pred to one-hot causing integer-index
errors for single-channel logits; update the UnifiedFocalLoss forward (or the
method handling y_pred) to either apply softmax/sigmoid before the one_hot
conversion or add an explicit input validation that raises a clear error when
y_pred has a single channel while use_softmax or use_sigmoid is True
(referencing variables y_pred, use_softmax, use_sigmoid and the one_hot
conversion block) so users get a helpful message rather than a silent failure.
Multi-class limitation: softmax won't work beyond binary despite being added.
The component losses (AsymmetricFocalLoss lines 135-138 and AsymmetricFocalTverskyLoss lines 79-80) hardcode indices [:, 0] and [:, 1], limiting support to exactly 2 classes. Even with use_softmax=True, inputs with >2 channels will fail or produce incorrect results.
Consider documenting this limitation in the use_softmax docstring or adding validation to reject num_classes != 2.
🤖 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/unified_focal_loss.py` around lines 241 - 244, The code enables
softmax via use_softmax but the component losses AsymmetricFocalLoss and
AsymmetricFocalTverskyLoss index hardcoded class channels ([:, 0] and [:, 1]) so
softmax with >2 channels will be broken; update the functions to validate and
reject multi-class inputs (e.g., check y_pred.shape[1] or num_classes and raise
a clear ValueError if != 2) or explicitly document in the use_softmax docstring
that only binary (2-class) predictions are supported, and ensure the error
message references use_softmax, AsymmetricFocalLoss, and
AsymmetricFocalTverskyLoss so callers know why their multi-channel input is
unsupported.
…@gmail.com> I, Oleksandr Yizchak Sanin <alexaaander.sanin@gmail.com>, hereby add my Signed-off-by to this commit: ec11cc3 Signed-off-by: Oleksandr Yizchak Sanin <alexaaander.sanin@gmail.com>
Move the sigmoid/softmax activation step before the one_hot conversion in AsymmetricUnifiedFocalLoss.forward(). The one_hot function uses input values as scatter indices, so passing raw logits (e.g. 10.0) causes "index out of bounds" errors. Activation must convert logits to probabilities first. Signed-off-by: Oleksandr Yizchak Sanin <alexaaander.sanin@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #8603
Adds `use_softmax` and `use_sigmoid` parameters to `AsymmetricUnifiedFocalLoss`, following the same pattern as `FocalLoss`. This allows users to pass raw logits directly without manually applying activations beforehand.
Test plan
Signed-off-by: Oleksandr Sanin alexaaander.sanin@gmail.com