diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py index a8d81eb579..e736851204 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 @@ -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: diff --git a/tests/integration/test_vis_gradbased.py b/tests/integration/test_vis_gradbased.py index e9db0af240..92dc5a8426 100644 --- a/tests/integration/test_vis_gradbased.py +++ b/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()