From 03dacb829504459ee605961ec6fbe23bdab99573 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 20 Aug 2026 22:26:12 +0100 Subject: [PATCH 1/6] perf(visualize): batch SmoothGrad noise samples with chunked forward Add sample_batch_size knob to SmoothGrad to evaluate n_samples noisy copies in batched chunks, amortizing kernel launches. Default 1 preserves exact current behavior and RNG stream. When chunk > 1 the class index is resolved once on the clean input so all copies use the same class, and BatchNorm-train-mode is warned. Signed-off-by: Soumya Snigdha Kundu --- monai/visualize/gradient_based.py | 66 +++++++++++++++++++++++++------ 1 file changed, 55 insertions(+), 11 deletions(-) diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py index a8d81eb5790..7b3491c98ec 100644 --- a/monai/visualize/gradient_based.py +++ b/monai/visualize/gradient_based.py @@ -11,6 +11,8 @@ from __future__ import annotations +import math +import warnings from collections.abc import Callable from functools import partial from typing import Any @@ -89,8 +91,6 @@ def model(self, m): def get_grad( self, x: torch.Tensor, index: torch.Tensor | int | None, retain_graph: bool = True, **kwargs: Any ) -> torch.Tensor: - if x.shape[0] != 1: - raise ValueError("expect batch size of 1") x.requires_grad = True self._model(x, class_idx=index, retain_graph=retain_graph, **kwargs) @@ -118,29 +118,73 @@ def __init__( n_samples: int = 25, magnitude: bool = True, verbose: bool = True, + sample_batch_size: int = 1, ) -> None: super().__init__(model) self.stdev_spread = stdev_spread self.n_samples = n_samples self.magnitude = magnitude + self.sample_batch_size = sample_batch_size + if self.sample_batch_size < 1: + raise ValueError("sample_batch_size must be >= 1") self.range: Callable if verbose and has_trange: self.range = partial(trange, desc=f"Computing {self.__class__.__name__}") else: self.range = range + def _resolve_index( + self, x: torch.Tensor, index: torch.Tensor | int | None, **kwargs: Any + ) -> torch.Tensor | int | None: + if index is not None: + return index + # resolve argmax once on clean input + was_training = self._model.model.training + self._model.model.eval() + with torch.no_grad(): + logits = self._model.model(x, **kwargs) + if isinstance(logits, (list, tuple)): + logits = logits[0] + if was_training: + self._model.model.train() + # logits shape (B, C) with B==1 for SmoothGrad input + resolved: int = int(logits.argmax(dim=1).item()) + return resolved + def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **kwargs: Any) -> torch.Tensor: + if self.sample_batch_size > 1 and self._model.model.training: + warnings.warn("SmoothGrad with sample_batch_size > 1 and model in train mode: BatchNorm statistics will mix noisy copies.") stdev = (self.stdev_spread * (x.max() - x.min())).item() total_gradients = torch.zeros_like(x) - for _ in self.range(self.n_samples): - # create noisy image - noise = torch.normal(0, stdev, size=x.shape, dtype=torch.float32, device=x.device) - x_plus_noise = x + noise - x_plus_noise = x_plus_noise.detach() - - # get gradient and accumulate - grad = self.get_grad(x_plus_noise, index, **kwargs) - total_gradients += (grad * grad) if self.magnitude else grad + # resolve index once so all noisy copies use the same class + resolved_index = self._resolve_index(x, index, **kwargs) if self.sample_batch_size > 1 else index + # if chunking, resolve None to a concrete int so batched forward uses consistent class + if resolved_index is None and self.sample_batch_size > 1: + # fallback: _resolve_index already handled None case above, but keep guard + resolved_index = self._resolve_index(x, None, **kwargs) + # for sample_batch_size==1, keep original per-sample argmax behaviour when index is None + # but when sample_batch_size>1 we must use resolved_index + n_chunks = math.ceil(self.n_samples / self.sample_batch_size) + remaining = self.n_samples + for _ in self.range(n_chunks): + k = min(self.sample_batch_size, remaining) + remaining -= k + if k == 1 and self.sample_batch_size == 1: + noise = torch.normal(0, stdev, size=x.shape, dtype=torch.float32, device=x.device) + x_plus_noise = (x + noise).detach() + grad = self.get_grad(x_plus_noise, index, **kwargs) + total_gradients += (grad * grad) if self.magnitude else grad + else: + # batched path + noise_shape = (k, *x.shape[1:]) + noise = torch.normal(0, stdev, size=noise_shape, dtype=torch.float32, device=x.device) + # x is (1, ...) -> broadcast to (k, ...) + x_plus_noise = (x + noise).detach() + grads = self.get_grad(x_plus_noise, resolved_index, **kwargs) + # grads shape (k, ...) + if self.magnitude: + grads = grads * grads + total_gradients += grads.sum(dim=0, keepdim=True) # average if self.magnitude: From 15a833d0e29b8afc65981c92c31a2f9e2354f511 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 20 Aug 2026 22:32:27 +0100 Subject: [PATCH 2/6] fix(visualize): reject non-integer sample_batch_size Validate that sample_batch_size is a real int before it reaches torch.normal as a shape dimension. Previously 1.5 passed the >= 1 check and produced a float size, and True passed because bool subclasses int. Both are now rejected at construction time. Addresses CodeRabbit review on #9064 (monai/visualize/gradient_based.py:129). Signed-off-by: Soumya Snigdha Kundu --- monai/visualize/gradient_based.py | 6 ++++-- tests/integration/test_vis_gradbased.py | 8 ++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py index 7b3491c98ec..3e3303510b2 100644 --- a/monai/visualize/gradient_based.py +++ b/monai/visualize/gradient_based.py @@ -124,9 +124,11 @@ def __init__( self.stdev_spread = stdev_spread self.n_samples = n_samples self.magnitude = magnitude + if isinstance(sample_batch_size, bool) or not isinstance(sample_batch_size, int): + raise ValueError(f"sample_batch_size must be an int, got {type(sample_batch_size).__name__}.") + if sample_batch_size < 1: + raise ValueError(f"sample_batch_size must be >= 1, got {sample_batch_size}.") self.sample_batch_size = sample_batch_size - if self.sample_batch_size < 1: - raise ValueError("sample_batch_size must be >= 1") self.range: Callable if verbose and has_trange: self.range = partial(trange, desc=f"Computing {self.__class__.__name__}") diff --git a/tests/integration/test_vis_gradbased.py b/tests/integration/test_vis_gradbased.py index e9db0af2404..aa4ffec44ff 100644 --- a/tests/integration/test_vis_gradbased.py +++ b/tests/integration/test_vis_gradbased.py @@ -65,5 +65,13 @@ def test_shape(self, vis_type, model, shape): self.assertTupleEqual(result.shape, x.shape) +class TestSmoothGradSampleBatchSize(unittest.TestCase): + + @parameterized.expand([[0], [-1], [1.5], [True], [False], ["2"]]) + def test_invalid_sample_batch_size(self, value): + with self.assertRaises(ValueError): + SmoothGrad(DENSENET2D, sample_batch_size=value) + + if __name__ == "__main__": unittest.main() From d5f9e4d30c0f1d20ecc2bc02d623df25571ef61f Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 20 Aug 2026 22:33:24 +0100 Subject: [PATCH 3/6] fix(visualize): restore every module mode after index resolution _resolve_index saved and restored only the root module's training flag, so a submodule deliberately left in eval mode was switched back on, and an exception in the clean forward left the whole model stuck in eval. Snapshot every submodule's mode and restore it in a finally block. Addresses CodeRabbit review on #9064 (monai/visualize/gradient_based.py:149). Signed-off-by: Soumya Snigdha Kundu --- monai/visualize/gradient_based.py | 21 ++++++++++--------- tests/integration/test_vis_gradbased.py | 27 +++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py index 3e3303510b2..dc39c00949e 100644 --- a/monai/visualize/gradient_based.py +++ b/monai/visualize/gradient_based.py @@ -140,15 +140,18 @@ def _resolve_index( ) -> torch.Tensor | int | None: if index is not None: return index - # resolve argmax once on clean input - was_training = self._model.model.training - self._model.model.eval() - with torch.no_grad(): - logits = self._model.model(x, **kwargs) - if isinstance(logits, (list, tuple)): - logits = logits[0] - if was_training: - self._model.model.train() + # resolve argmax once on clean input, restoring every submodule's mode afterwards + modules = tuple(self._model.model.modules()) + training_states = [m.training for m in modules] + try: + self._model.model.eval() + with torch.no_grad(): + logits = self._model.model(x, **kwargs) + if isinstance(logits, (list, tuple)): + logits = logits[0] + finally: + for m, training in zip(modules, training_states): + m.train(training) # logits shape (B, C) with B==1 for SmoothGrad input resolved: int = int(logits.argmax(dim=1).item()) return resolved diff --git a/tests/integration/test_vis_gradbased.py b/tests/integration/test_vis_gradbased.py index aa4ffec44ff..fe65d4319a4 100644 --- a/tests/integration/test_vis_gradbased.py +++ b/tests/integration/test_vis_gradbased.py @@ -73,5 +73,32 @@ def test_invalid_sample_batch_size(self, value): SmoothGrad(DENSENET2D, sample_batch_size=value) +class TestSmoothGradModelState(unittest.TestCase): + + def test_mixed_module_modes_are_restored(self): + model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3) + model.train() + # a submodule deliberately kept in eval mode must stay there + frozen = next(m for m in model.modules() if isinstance(m, torch.nn.BatchNorm2d)) + frozen.eval() + before = [m.training for m in model.modules()] + + vis = SmoothGrad(model, n_samples=2, sample_batch_size=2, verbose=False) + vis._resolve_index(torch.rand(1, 1, 48, 64), None) + + self.assertEqual([m.training for m in model.modules()], before) + + def test_mode_restored_when_forward_raises(self): + model = DenseNetAdjoint(spatial_dims=2, in_channels=1, out_channels=3) + model.train() + before = [m.training for m in model.modules()] + + vis = SmoothGrad(model, n_samples=2, sample_batch_size=2, verbose=False) + with self.assertRaises(ValueError): + vis._resolve_index(torch.rand(1, 1, 48, 64), None, adjoint_info=0) + + self.assertEqual([m.training for m in model.modules()], before) + + if __name__ == "__main__": unittest.main() From b5006313f7ec76a850f0192994686482a7f95fe8 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 20 Aug 2026 22:33:55 +0100 Subject: [PATCH 4/6] fix(visualize): restore the batch-size-1 guard on VanillaGrad.get_grad The guard was dropped outright so the batched SmoothGrad path could call get_grad with a chunk of noisy copies. That also silently relaxed the contract of VanillaGrad and GuidedBackpropGrad, whose public entry points would sum gradients over the batch instead of raising. Keep the guard on get_grad and route the batched path through a private _get_grad, and reject input batches > 1 in SmoothGrad when sample_batch_size > 1, where x + noise broadcasting and the sum(dim=0) reduction assume a single clean input. Addresses CodeRabbit review on #9064 (monai/visualize/gradient_based.py:187). Signed-off-by: Soumya Snigdha Kundu --- monai/visualize/gradient_based.py | 12 +++++++++++- tests/integration/test_vis_gradbased.py | 14 ++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py index dc39c00949e..7a9e0d20982 100644 --- a/monai/visualize/gradient_based.py +++ b/monai/visualize/gradient_based.py @@ -91,6 +91,14 @@ def model(self, m): def get_grad( self, x: torch.Tensor, index: torch.Tensor | int | None, retain_graph: bool = True, **kwargs: Any ) -> torch.Tensor: + if x.shape[0] != 1: + raise ValueError("expect batch size of 1") + return self._get_grad(x, index, retain_graph=retain_graph, **kwargs) + + def _get_grad( + self, x: torch.Tensor, index: torch.Tensor | int | None, retain_graph: bool = True, **kwargs: Any + ) -> torch.Tensor: + # unguarded gradient computation; ``x`` may carry a batch of noisy copies of a single input x.requires_grad = True self._model(x, class_idx=index, retain_graph=retain_graph, **kwargs) @@ -157,6 +165,8 @@ def _resolve_index( return resolved def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **kwargs: Any) -> torch.Tensor: + if self.sample_batch_size > 1 and x.shape[0] != 1: + raise ValueError(f"sample_batch_size > 1 expects an input batch size of 1, got {x.shape[0]}.") if self.sample_batch_size > 1 and self._model.model.training: warnings.warn("SmoothGrad with sample_batch_size > 1 and model in train mode: BatchNorm statistics will mix noisy copies.") stdev = (self.stdev_spread * (x.max() - x.min())).item() @@ -185,7 +195,7 @@ def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **k noise = torch.normal(0, stdev, size=noise_shape, dtype=torch.float32, device=x.device) # x is (1, ...) -> broadcast to (k, ...) x_plus_noise = (x + noise).detach() - grads = self.get_grad(x_plus_noise, resolved_index, **kwargs) + grads = self._get_grad(x_plus_noise, resolved_index, **kwargs) # grads shape (k, ...) if self.magnitude: grads = grads * grads diff --git a/tests/integration/test_vis_gradbased.py b/tests/integration/test_vis_gradbased.py index fe65d4319a4..1f66131f364 100644 --- a/tests/integration/test_vis_gradbased.py +++ b/tests/integration/test_vis_gradbased.py @@ -100,5 +100,19 @@ def test_mode_restored_when_forward_raises(self): self.assertEqual([m.training for m in model.modules()], before) +class TestGradBatchContract(unittest.TestCase): + + @parameterized.expand([[VanillaGrad], [GuidedBackpropGrad]]) + def test_get_grad_rejects_batch(self, vis_type): + vis = vis_type(DENSENET2D) + with self.assertRaisesRegex(ValueError, "batch size of 1"): + vis(torch.rand(2, 1, 48, 64)) + + def test_smoothgrad_batched_rejects_input_batch(self): + vis = SmoothGrad(DENSENET2D, n_samples=2, sample_batch_size=2, verbose=False) + with self.assertRaisesRegex(ValueError, "input batch size of 1"): + vis(torch.rand(2, 1, 48, 64)) + + if __name__ == "__main__": unittest.main() From 8cee0f690d246e599fba63cc7dfc2272e12501af Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 20 Aug 2026 22:34:22 +0100 Subject: [PATCH 5/6] fix(visualize): normalise a tensor class index for the batched path _resolve_index returned a caller-supplied index unchanged, so a tensor index shaped for the clean input was fed to a forward of size k and selected the wrong class per noisy copy. Convert a single-element tensor to an int and reject multi-element ones, which cannot describe one shared class. Also covers batched-vs-unbatched equivalence for a fixed seed. Signed-off-by: Soumya Snigdha Kundu --- monai/visualize/gradient_based.py | 8 ++++++++ tests/integration/test_vis_gradbased.py | 24 ++++++++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py index 7a9e0d20982..810746bd70f 100644 --- a/monai/visualize/gradient_based.py +++ b/monai/visualize/gradient_based.py @@ -147,6 +147,14 @@ def _resolve_index( self, x: torch.Tensor, index: torch.Tensor | int | None, **kwargs: Any ) -> torch.Tensor | int | None: if index is not None: + if isinstance(index, torch.Tensor): + # the batched path feeds a forward of size k, so a per-sample index tensor + # sized for the clean input would select the wrong class per copy + if index.numel() != 1: + raise ValueError( + f"sample_batch_size > 1 expects a single class index, got {index.numel()} elements." + ) + return int(index.item()) return index # resolve argmax once on clean input, restoring every submodule's mode afterwards modules = tuple(self._model.model.modules()) diff --git a/tests/integration/test_vis_gradbased.py b/tests/integration/test_vis_gradbased.py index 1f66131f364..92dc5a8426f 100644 --- a/tests/integration/test_vis_gradbased.py +++ b/tests/integration/test_vis_gradbased.py @@ -114,5 +114,29 @@ def test_smoothgrad_batched_rejects_input_batch(self): vis(torch.rand(2, 1, 48, 64)) +class TestSmoothGradIndex(unittest.TestCase): + + def test_tensor_index_is_normalised(self): + vis = SmoothGrad(DENSENET2D, n_samples=2, sample_batch_size=2, verbose=False) + self.assertEqual(vis._resolve_index(torch.rand(1, 1, 48, 64), torch.tensor([2])), 2) + + def test_multi_element_tensor_index_rejected(self): + vis = SmoothGrad(DENSENET2D, n_samples=2, sample_batch_size=2, verbose=False) + with self.assertRaisesRegex(ValueError, "single class index"): + vis._resolve_index(torch.rand(1, 1, 48, 64), torch.tensor([0, 1])) + + def test_batched_matches_unbatched(self): + model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3).eval() + x = torch.rand(1, 1, 48, 64) + + torch.manual_seed(0) + expected = SmoothGrad(model, n_samples=4, sample_batch_size=1, verbose=False)(x, index=1) + torch.manual_seed(0) + actual = SmoothGrad(model, n_samples=4, sample_batch_size=4, verbose=False)(x, index=1) + + self.assertTupleEqual(actual.shape, x.shape) + torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-3) + + if __name__ == "__main__": unittest.main() From b7cacceb75016fab50341f0c465443013eb16cbe Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 20 Aug 2026 22:36:25 +0100 Subject: [PATCH 6/6] refactor(visualize): drop unreachable index fallback, fix warn lint _resolve_index never returns None once index resolution runs, so the "resolved_index is None" fallback and its stale comment were dead. Also give the train-mode warning a stacklevel and wrap it to satisfy flake8 B028/E501. Signed-off-by: Soumya Snigdha Kundu --- monai/visualize/gradient_based.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py index 810746bd70f..e7368512042 100644 --- a/monai/visualize/gradient_based.py +++ b/monai/visualize/gradient_based.py @@ -176,17 +176,15 @@ def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **k if self.sample_batch_size > 1 and x.shape[0] != 1: raise ValueError(f"sample_batch_size > 1 expects an input batch size of 1, got {x.shape[0]}.") if self.sample_batch_size > 1 and self._model.model.training: - warnings.warn("SmoothGrad with sample_batch_size > 1 and model in train mode: BatchNorm statistics will mix noisy copies.") + warnings.warn( + "SmoothGrad with sample_batch_size > 1 and model in train mode: " + "BatchNorm statistics will mix noisy copies.", + stacklevel=2, + ) stdev = (self.stdev_spread * (x.max() - x.min())).item() total_gradients = torch.zeros_like(x) # resolve index once so all noisy copies use the same class resolved_index = self._resolve_index(x, index, **kwargs) if self.sample_batch_size > 1 else index - # if chunking, resolve None to a concrete int so batched forward uses consistent class - if resolved_index is None and self.sample_batch_size > 1: - # fallback: _resolve_index already handled None case above, but keep guard - resolved_index = self._resolve_index(x, None, **kwargs) - # for sample_batch_size==1, keep original per-sample argmax behaviour when index is None - # but when sample_batch_size>1 we must use resolved_index n_chunks = math.ceil(self.n_samples / self.sample_batch_size) remaining = self.n_samples for _ in self.range(n_chunks):