Use binary_thinning_3d_cuda in SoftclDiceLoss#8904
Conversation
📝 WalkthroughWalkthroughAdds an optional CUDA hard-target skeletonization path to clDice losses: conditional import of binary_thinning_3d, a use_hard_target constructor flag on SoftclDiceLoss and SoftDiceclDiceLoss, runtime validation and per-batch/channel thinning when enabled, dependency extras in setup.cfg for binary_thinning_3d_cuda, and a CUDA-gated unit test for the hard-target path. Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (1)
monai/losses/cldice.py (1)
321-322:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winClarify dependency name in docstring.
Same issue as SoftclDiceLoss: docstring mentions "MONAI C++ extensions" but should specify
binary_thinning_3d_cuda.🤖 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/cldice.py` around lines 321 - 322, Update the docstring for the use_hard_target parameter to explicitly name the required extension: replace the vague "MONAI C++ extensions" text with a clear reference to `binary_thinning_3d_cuda` (the MONAI C++ extension providing 3D binary thinning); ensure the description on use_hard_target in clDiceLoss (and mirror the same wording used for SoftclDiceLoss) states that `binary_thinning_3d_cuda` is required for CUDA 3D target thinning and that this option defaults to False.
🧹 Nitpick comments (2)
monai/losses/cldice.py (1)
239-239: ⚡ Quick winDocument the second argument to binary_thinning.
The second argument
0tobinary_thinning_3d.binary_thinningis undocumented. Add a comment explaining its purpose.🤖 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/cldice.py` at line 239, Add a short inline comment explaining the purpose of the magic numeric second argument passed to binary_thinning_3d.binary_thinning(skel_true[b, c], 0); specifically state what the value 0 represents (e.g., the connectivity/foreground label/iteration mode used by binary_thinning) and why that choice is appropriate here; locate the call in cldice.py where binary_thinning_3d.binary_thinning is invoked (on skel_true and skel_pred slices) and append a concise comment clarifying the semantic meaning of the second parameter and its expected values.tests/losses/test_cldice_loss.py (1)
88-93: 🏗️ Heavy liftExpand test coverage for hard target mode.
Test only covers trivial case where input equals target. Add test cases with imperfect overlap to verify hard skeletonization produces expected loss values compared to soft mode.
🤖 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_cldice_loss.py` around lines 88 - 93, Add additional assertions to test_hard_target in tests/losses/test_cldice_loss.py: create non-trivial 3D examples where prediction and target have imperfect overlap (e.g., single-voxel shifts or partial overlap using the same ONES_3D shape convention), compute loss using SoftclDiceLoss(use_hard_target=True) and compare against SoftclDiceLoss(use_hard_target=False) (or known numeric expectations) to verify hard skeletonization changes the value; ensure calls use the loss(...) invocation on CUDA tensors (as with ONES_3D["input"].cuda()) and assert the hard-mode loss is different and matches the expected relation to soft-mode (e.g., greater or lower depending on your algorithm) with np.testing.assert_allclose or assert_not_equal as appropriate.Source: Coding guidelines
🤖 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/cldice.py`:
- Around line 157-158: The docstring for the use_hard_target parameter in
monai.losses.cldice.py is ambiguous about the required dependency; update the
parameter description (the use_hard_target docstring in the cLDice/CLDice loss
implementation) to explicitly name the dependency package
`binary_thinning_3d_cuda` and add a brief install hint (e.g., "install
binary_thinning_3d_cuda to enable CUDA 3D binary thinning"). Keep the existing
note about requiring a 3D target and MONAI C++ extensions if still relevant, but
ensure the exact package name `binary_thinning_3d_cuda` is present so users can
install the correct dependency.
- Around line 235-242: When use_hard_target is True but the code falls back to
soft_skel (because target.dim() != 5, _has_thinning is False, or target.is_cuda
is False), emit a clear warning instead of silently continuing; locate the
branch that computes skel_true (the if using binary_thinning_3d.binary_thinning
and the else calling soft_skel) and add a warnings.warn (or use the module
logger) indicating that hard-target thinning was requested but unavailable and
that soft_skel is being used as fallback, including the reason (which condition
failed) and referencing use_hard_target, binary_thinning_3d.binary_thinning, and
soft_skel in the message.
In `@tests/losses/test_cldice_loss.py`:
- Around line 88-93: The test test_hard_target currently may pass via the soft
fallback if binary_thinning_3d is unavailable; modify the test to first check
for the binary_thinning_3d dependency and skip the test when it's not
importable, or explicitly verify the hard-target branch is used by asserting
that SoftclDiceLoss(use_hard_target=True) actually calls/relies on
binary_thinning_3d (e.g., by checking importability or observing a
side-effect/flag from the hard-skeletonization path). Reference SoftclDiceLoss
and binary_thinning_3d in the test_hard_target test and either skip when
binary_thinning_3d is missing or add an assertion that confirms the hard-target
implementation was executed.
- Line 89: Add a PEP 257 docstring to the test_hard_target method in
tests/losses/test_cldice_loss.py: open the test_hard_target function and add a
one-line or short multi-line docstring that succinctly describes the behavior
being tested (e.g., what inputs are provided, expected outcome, and any edge
conditions like hard targets), ensuring it follows project docstring style
guidelines.
---
Duplicate comments:
In `@monai/losses/cldice.py`:
- Around line 321-322: Update the docstring for the use_hard_target parameter to
explicitly name the required extension: replace the vague "MONAI C++ extensions"
text with a clear reference to `binary_thinning_3d_cuda` (the MONAI C++
extension providing 3D binary thinning); ensure the description on
use_hard_target in clDiceLoss (and mirror the same wording used for
SoftclDiceLoss) states that `binary_thinning_3d_cuda` is required for CUDA 3D
target thinning and that this option defaults to False.
---
Nitpick comments:
In `@monai/losses/cldice.py`:
- Line 239: Add a short inline comment explaining the purpose of the magic
numeric second argument passed to
binary_thinning_3d.binary_thinning(skel_true[b, c], 0); specifically state what
the value 0 represents (e.g., the connectivity/foreground label/iteration mode
used by binary_thinning) and why that choice is appropriate here; locate the
call in cldice.py where binary_thinning_3d.binary_thinning is invoked (on
skel_true and skel_pred slices) and append a concise comment clarifying the
semantic meaning of the second parameter and its expected values.
In `@tests/losses/test_cldice_loss.py`:
- Around line 88-93: Add additional assertions to test_hard_target in
tests/losses/test_cldice_loss.py: create non-trivial 3D examples where
prediction and target have imperfect overlap (e.g., single-voxel shifts or
partial overlap using the same ONES_3D shape convention), compute loss using
SoftclDiceLoss(use_hard_target=True) and compare against
SoftclDiceLoss(use_hard_target=False) (or known numeric expectations) to verify
hard skeletonization changes the value; ensure calls use the loss(...)
invocation on CUDA tensors (as with ONES_3D["input"].cuda()) and assert the
hard-mode loss is different and matches the expected relation to soft-mode
(e.g., greater or lower depending on your algorithm) with
np.testing.assert_allclose or assert_not_equal as appropriate.
🪄 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: bd9b7761-2e68-4412-acb5-0622d96383e4
📒 Files selected for processing (3)
monai/losses/cldice.pysetup.cfgtests/losses/test_cldice_loss.py
Signed-off-by: yang <sychen52@gmail.com>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
monai/losses/cldice.py (1)
193-202: ⚡ Quick winDocument
ValueErrorforuse_hard_targetconditions in Raises section.The forward method now raises
ValueErrorwhenuse_hard_target=Truebut conditions aren't met (lines 237-240). Add this to the docstring's Raises section.📝 Suggested docstring update
Raises: AssertionError: When input and target (after one hot transform if set) have different shapes. + ValueError: When ``use_hard_target=True`` but target is not a 5D CUDA tensor + or binary_thinning_3d_cuda is unavailable. """🤖 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/cldice.py` around lines 193 - 202, Update the forward method docstring in cldice.py (function forward) to document the new ValueError raised when use_hard_target=True and the target does not meet required conditions; specifically add a Raises entry that states a ValueError is raised if use_hard_target is True but the target tensor is not binary/hard (or not of expected shape/type) so the hard-target conversion cannot be applied. Reference the forward method and the use_hard_target branch where the ValueError is thrown and ensure the Raises section lists both AssertionError (existing) and this ValueError.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@monai/losses/cldice.py`:
- Around line 193-202: Update the forward method docstring in cldice.py
(function forward) to document the new ValueError raised when
use_hard_target=True and the target does not meet required conditions;
specifically add a Raises entry that states a ValueError is raised if
use_hard_target is True but the target tensor is not binary/hard (or not of
expected shape/type) so the hard-target conversion cannot be applied. Reference
the forward method and the use_hard_target branch where the ValueError is thrown
and ensure the Raises section lists both AssertionError (existing) and this
ValueError.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: d9880156-2a83-4da0-becc-6ee551f1032d
📒 Files selected for processing (3)
monai/losses/cldice.pysetup.cfgtests/losses/test_cldice_loss.py
🚧 Files skipped from review as they are similar to previous changes (2)
- setup.cfg
- tests/losses/test_cldice_loss.py
Description
This PR introduces the
use_hard_targetparameter toSoftclDiceLossandSoftDiceclDiceLoss, allowing users to optionally compute the ground-truth topological skeleton using an exact, CUDA-accelerated 3D binary thinning algorithm.Currently, the clDice loss relies on
soft_skel(morphological max-pooling) for both the prediction and the target. Whilesoft_skelis necessary for the prediction to maintain differentiability, it provides a thick, fuzzy approximation of a centerline.By dynamically importing the
binary_thinning_3d_cudapackage for the target mask, this PR brings two major advantages:Mathematical Accuracy: It evaluates the topological sensitivity (
tsens) against an exact, 1-pixel-wide, topologically-correct hard centerline, remaining truer to the original clDice paper's formulation.High Performance: Because the algorithm is fully implemented in CUDA, it extracts the 3D skeleton in milliseconds directly on the GPU, avoiding CPU bottlenecks during the training loop.
Mathematical Accuracy: It evaluates the topological sensitivity ( tsens ) against an exact, 1-
pixel-wide, topologically-correct hard centerline, remaining truer to the original clDice paper's
formulation.
High Performance: Because the algorithm is fully implemented in CUDA, it extracts the 3D skeleton
in milliseconds directly on the GPU, avoiding CPU bottlenecks during the training loop.
Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.