perf(visualize): batch SmoothGrad noise samples with chunked forward

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
@@ -91,6 +93,12 @@ def get_grad(
     ) -> 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)
@@ -118,29 +126,86 @@ 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
+        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
         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:
+            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())
+        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
+
     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.",
+                stacklevel=2,
+            )
         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
+        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:

tests/integration/test_vis_gradbased.py

@@ -65,5 +65,78 @@ 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)
+
+
+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)
+
+
+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))
+
+
+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()