Add batch support for saliency maps
Fixes #7334 .
Description
This PR adds batch processing support for saliency methods (e.g., VanillaGrad) in MONAI.
Types of changes
- [x] Non-breaking change (fix or new feature that would not break existing functionality).
- [x] New tests added to cover the changes.
- [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.
Walkthrough
- monai/visualize/class_activation_maps.py: Extended ModelWithHooks.class_score to accept int or torch.Tensor, including per-batch class index selection via torch.gather. Updated call to handle Tensor class_idx (moving to logits device) and to pass class_idx directly to class_score without int casting. Forward/backward hook flow unchanged.
- monai/visualize/gradient_based.py: Removed batch-size==1 validation in VanillaGrad.get_grad, enabling gradients for batched inputs; computation path unchanged otherwise.
- tests/integration/test_vis_gradbased.py: Added parameterized cases with larger/batched input shapes across multiple model types to cover batch processing paths.
Estimated code review effort
🎯 3 (Moderate) | ⏱️ ~25 minutes
Pre-merge checks and finishing touches
❌ Failed checks (1 warning, 1 inconclusive)
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage | ⚠️ Warning | Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. | You can run @coderabbitai generate docstrings to improve docstring coverage. |
| Out of Scope Changes Check | ❓ Inconclusive | Most changes are directly related to enabling batch saliency support (gradient_based.py, ModelWithHooks, and expanded tests). However, the raw summary contains a contradictory note about ModelWithHooks.class_score's declaration (the narrative says it was extended to accept int | Tensor but the "Alterations" line indicates a signature of class_idx: int), creating ambiguity about a possible unintended public API change. Because of this inconsistency in the provided summary, a conclusive determination about any breaking or out-of-scope public API modification cannot be made. |
✅ Passed checks (3 passed)
| Check name | Status | Explanation |
|---|---|---|
| Title Check | ✅ Passed | The title "Add batch support for saliency maps" is concise, focused, and accurately summarizes the primary intent of the PR—enabling batch processing for saliency methods (e.g., removing single-sample enforcement and updating scoring behavior). It is clear and free of noise, so teammates scanning history will understand the main change. |
| Linked Issues Check | ✅ Passed | The changes address issue #7334 by removing the single-sample restriction in VanillaGrad.get_grad, enabling per-sample class_idx handling in ModelWithHooks, and adding parameterized tests that exercise multiple batch shapes and gradient-based methods. These modifications implement the requested batch-processing capability for saliency map generation across multiple input samples. Based on the provided summaries, the implementation aligns with the linked issue objectives. |
| Description Check | ✅ Passed | The PR description follows the repository template: it includes "Fixes #7334", a short Description of the change (adding batch processing support for saliency methods), and the "Types of changes" checklist with tests and local runs marked as passed. The author also states integration and quick tests passed locally. The description is sufficient for review but could be slightly improved with a brief implementation note or list of modified files. |
✨ Finishing touches
- [ ] 📝 Generate Docstrings
🧪 Generate unit tests
- [ ] Create PR with unit tests
- [ ] Post copyable unit tests in a comment
[!TIP]
👮 Agentic pre-merge checks are now available in preview!
Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
- Built-in checks – Quickly apply ready-made checks to enforce title conventions, require pull request descriptions that follow templates, validate linked issues for compliance, and more.
- Custom agentic checks – Define your own rules using CodeRabbit’s advanced agentic capabilities to enforce organization-specific policies and workflows. For example, you can instruct CodeRabbit’s agent to verify that API documentation is updated whenever API schema files are modified in a PR. Note: Upto 5 custom checks are currently allowed during the preview period. Pricing for this feature will be announced in a few weeks.
Please see the documentation for more information.
Example:
reviews: pre_merge_checks: custom_checks: - name: "Undocumented Breaking Changes" mode: "warning" instructions: | Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post.
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 help to get the list of available commands and usage tips.
Hi @einsyang723 thanks for the contribution. There are some minor issues with typing to address and the comments from Coderabbit, if these can be addressed we can review again and consider merging.