Support batched inputs in gradient-based saliency maps

monai/visualize/class_activation_maps.py

@@ -73,7 +73,7 @@ def __init__(
         self.gradients: dict[str, torch.Tensor] = {}
         self.activations: dict[str, torch.Tensor] = {}
         self.score: torch.Tensor | None = None
-        self.class_idx: int | None = None
+        self.class_idx: int | torch.Tensor | None = None
         self.register_backward = register_backward
         self.register_forward = register_forward
 
@@ -124,29 +124,54 @@ def get_layer(self, layer_id: str | Callable[[nn.Module], nn.Module]) -> nn.Modu
                     return cast(nn.Module, mod)
         raise NotImplementedError(f"Could not find {layer_id}.")
 
-    def class_score(self, logits: torch.Tensor, class_idx: int) -> torch.Tensor:
-        return logits[:, class_idx].squeeze()
+    def class_score(self, logits: torch.Tensor, class_idx: int | torch.Tensor) -> torch.Tensor:
+        """
+        Select the score of the requested class for every batch element.
+
+        Args:
+            logits: model output of shape ``(B, C, ...)``.
+            class_idx: an integer (the same class for every batch element), or a tensor with one class index
+                per batch element, of shape ``(B,)`` or ``(B, ...)`` matching the trailing dimensions of ``logits``.
+
+        Returns:
+            a tensor of shape ``(B, ...)`` holding one score per batch element.
+        """
+        if isinstance(class_idx, int):
+            return logits[:, class_idx]
+        idx = torch.as_tensor(class_idx, dtype=torch.long, device=logits.device)
+        if idx.ndim == 0:
+            return logits[:, int(idx)]
+        if idx.shape[0] != logits.shape[0]:
+            raise ValueError(
+                f"expected one class index per batch element, got {idx.shape[0]} for batch size {logits.shape[0]}."
+            )
+        if idx.ndim == 1:  # one class per batch element, applied to all trailing dimensions of logits
+            return logits[torch.arange(logits.shape[0], device=logits.device), idx]
+        return logits.gather(1, idx.unsqueeze(1)).squeeze(1)  # one class per batch element and location
 
     def __call__(self, x, class_idx=None, retain_graph=False, **kwargs):
         train = self.model.training
         self.model.eval()
-        logits = self.model(x, **kwargs)
-        self.class_idx = logits.max(1)[-1] if class_idx is None else class_idx
-        acti, grad = None, None
-        if self.register_forward:
-            acti = tuple(self.activations[layer] for layer in self.target_layers)
-        if self.register_backward:
-            self.score = self.class_score(logits, cast(int, self.class_idx))
-            self.model.zero_grad()
-            self.score.sum().backward(retain_graph=retain_graph)
-            for layer in self.target_layers:
-                if layer not in self.gradients:
-                    warnings.warn(
-                        f"Backward hook for {layer} is not triggered; `requires_grad` of {layer} should be `True`."
-                    )
-            grad = tuple(self.gradients[layer] for layer in self.target_layers if layer in self.gradients)
-        if train:
-            self.model.train()
+        try:
+            logits = self.model(x, **kwargs)
+            class_idx = logits.max(1)[-1] if class_idx is None else class_idx
+            self.class_idx = class_idx
+            acti, grad = None, None
+            if self.register_forward:
+                acti = tuple(self.activations[layer] for layer in self.target_layers)
+            if self.register_backward:
+                self.score = self.class_score(logits, class_idx)
+                self.model.zero_grad()
+                self.score.sum().backward(retain_graph=retain_graph)
+                for layer in self.target_layers:
+                    if layer not in self.gradients:
+                        warnings.warn(
+                            f"Backward hook for {layer} is not triggered; `requires_grad` of {layer} should be `True`."
+                        )
+                grad = tuple(self.gradients[layer] for layer in self.target_layers if layer in self.gradients)
+        finally:  # restore the training mode also when the forward or backward pass raises
+            if train:
+                self.model.train()
         return logits, acti, grad
 
     def get_wrapped_net(self):
@@ -406,7 +431,11 @@ def compute_map(self, x, class_idx=None, retain_graph=False, layer_idx=-1, **kwa
         alpha_dr = alpha_nr.mul(2) + acti.mul(grad.pow(3)).view(b, c, -1).sum(-1).view(b, c, *[1] * len(spatial))
         alpha_dr = torch.where(alpha_dr != 0.0, alpha_dr, torch.ones_like(alpha_dr))
         alpha = alpha_nr.div(alpha_dr + 1e-7)
-        relu_grad = F.relu(cast(torch.Tensor, self.nn_module.score).exp() * grad)
+        # one class score per batch element, broadcast over the channel and spatial dimensions of its gradient;
+        # spatial logits give one score per location, averaged here: exp(score) only scales the map of each image
+        # by a positive factor, and the mean keeps that factor finite where a sum over many locations overflows
+        score = cast(torch.Tensor, self.nn_module.score).reshape(b, -1).mean(1).reshape(b, *([1] * (grad.ndim - 1)))
+        relu_grad = F.relu(score.exp() * grad)
         weights = (alpha * relu_grad).view(b, c, -1).sum(-1).view(b, c, *[1] * len(spatial))
         acti_map = (weights * acti).sum(1, keepdim=True)
         return F.relu(acti_map)

monai/visualize/gradient_based.py

@@ -63,6 +63,10 @@ class VanillaGrad:
 
     If ``index`` is None, argmax of the output logits will be used.
 
+    ``x`` may hold a batch of images: ``index`` is then either an integer applied to every image, or a tensor with
+    one class index per image. The model is run in evaluation mode and the selected class scores are summed, so
+    for networks without cross-batch interactions the saliency map of each image is the one it would get on its own.
+
     See also:
 
         - Simonyan et al. Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps
@@ -89,8 +93,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)
@@ -104,7 +106,8 @@ def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **k
 class SmoothGrad(VanillaGrad):
     """
     Compute averaged sensitivity map based on ``n_samples`` (Gaussian additive) of noisy versions
-    of the input image ``x``.
+    of the input image ``x``. The noise standard deviation is ``stdev_spread`` times the intensity range
+    of each image in the batch ``x``.
 
     See also:
 
@@ -130,11 +133,14 @@ def __init__(
             self.range = range
 
     def __call__(self, x: torch.Tensor, index: torch.Tensor | int | None = None, **kwargs: Any) -> torch.Tensor:
-        stdev = (self.stdev_spread * (x.max() - x.min())).item()
+        # noise scale of each batch element is ``stdev_spread`` times the intensity range of that element
+        x_flat = x.detach().reshape(x.shape[0], -1)
+        stdev = self.stdev_spread * (x_flat.max(dim=1)[0] - x_flat.min(dim=1)[0])
+        stdev = stdev.view(-1, *([1] * (x.ndim - 1)))
         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)
+            noise = torch.randn(x.shape, dtype=torch.float32, device=x.device) * stdev
             x_plus_noise = x + noise
             x_plus_noise = x_plus_noise.detach()
 

tests/integration/test_vis_gradbased.py

@@ -18,6 +18,7 @@
 
 from monai.networks.nets import DenseNet, DenseNet121, SEResNet50
 from monai.visualize import GuidedBackpropGrad, GuidedBackpropSmoothGrad, SmoothGrad, VanillaGrad
+from tests.test_utils import assert_allclose
 
 
 class DenseNetAdjoint(DenseNet121):
@@ -28,6 +29,14 @@ def __call__(self, x, adjoint_info):
         return super().__call__(x)
 
 
+class HalfSquaredNorm(torch.nn.Module):
+    """Logits ``(s, -s)`` with ``s = 0.5 * ||x||^2`` per image, so the input gradient of class 0 is the input."""
+
+    def forward(self, x):
+        s = 0.5 * x.pow(2).flatten(1).sum(1, keepdim=True)
+        return torch.cat([s, -s], dim=1)
+
+
 DENSENET2D = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3)
 DENSENET3D = DenseNet(spatial_dims=3, in_channels=1, out_channels=3, init_features=2, growth_rate=2, block_config=(6,))
 SENET2D = SEResNet50(spatial_dims=2, in_channels=3, num_classes=4)
@@ -38,14 +47,20 @@ def __call__(self, x, adjoint_info):
 for type in (VanillaGrad, SmoothGrad, GuidedBackpropGrad, GuidedBackpropSmoothGrad):
     # 2D densenet
     TESTS.append([type, DENSENET2D, (1, 1, 48, 64)])
+    TESTS.append([type, DENSENET2D, (2, 1, 48, 64)])
     # 3D densenet
     TESTS.append([type, DENSENET3D, (1, 1, 6, 6, 6)])
+    TESTS.append([type, DENSENET3D, (3, 1, 6, 6, 6)])
     # 2D senet
     TESTS.append([type, SENET2D, (1, 3, 64, 64)])
     # 3D senet
     TESTS.append([type, SENET3D, (1, 3, 8, 8, 48)])
     # 2D densenet - adjoint
     TESTS.append([type, DENSENET2DADJOINT, (1, 1, 48, 64)])
+    TESTS.append([type, DENSENET2DADJOINT, (2, 1, 48, 64)])
+
+# deterministic methods: a batch must give, image by image, the maps each image gets on its own
+TESTS_BATCH = [[VanillaGrad], [GuidedBackpropGrad]]
 
 
 class TestGradientClassActivationMap(unittest.TestCase):
@@ -64,6 +79,60 @@ def test_shape(self, vis_type, model, shape):
         result = vis(x, **kwargs)
         self.assertTupleEqual(result.shape, x.shape)
 
+    @parameterized.expand(TESTS_BATCH)
+    def test_batch_matches_single_sample(self, vis_type):
+        """
+        The saliency maps of a batch equal the maps computed one image at a time, for ``index`` given as
+        ``None`` (per-image argmax), an integer (same class for every image) and a tensor (one class per image).
+        """
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        model = DENSENET3D.to(device)
+        model.eval()
+        vis = vis_type(model)
+        x = torch.rand((3, 1, 6, 6, 6), device=device)
+        for index in (None, 1, torch.tensor([0, 2, 1], device=device)):
+            batch = vis(x.clone(), index=index)
+            single = []
+            for i in range(x.shape[0]):
+                index_i = int(index[i]) if isinstance(index, torch.Tensor) else index
+                single.append(vis(x[i : i + 1].clone(), index=index_i))
+            assert_allclose(batch, torch.cat(single), rtol=1e-4, atol=1e-5)
+
+    def test_smoothgrad_batch_without_noise(self):
+        """With zero noise and a single sample, ``SmoothGrad`` on a batch reduces to ``VanillaGrad``."""
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        model = DENSENET3D.to(device)
+        model.eval()
+        x = torch.rand((2, 1, 6, 6, 6), device=device)
+        smooth = SmoothGrad(model, stdev_spread=0.0, n_samples=1, magnitude=False, verbose=False)
+        assert_allclose(smooth(x.clone()), VanillaGrad(model)(x.clone()), rtol=1e-4, atol=1e-5)
+
+    def test_smoothgrad_noise_per_image(self):
+        """
+        ``SmoothGrad`` scales the noise of each image with the intensity range of that image: with a model whose
+        input gradient is the input, one sample returns ``x + noise``, so the noise of an image with range about 1
+        must not take the scale of an image with range about 100, as it would with one range for the whole batch.
+        """
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        torch.manual_seed(0)
+        x = torch.rand((2, 1, 16, 16, 16), device=device)
+        x[1] *= 100.0
+        stdev_spread = 0.15
+        smooth = SmoothGrad(HalfSquaredNorm(), stdev_spread=stdev_spread, n_samples=1, magnitude=False, verbose=False)
+        noise = smooth(x.clone(), index=0) - x
+        for i in range(x.shape[0]):
+            expected = stdev_spread * (x[i].max() - x[i].min()).item()
+            self.assertAlmostEqual(noise[i].std().item() / expected, 1.0, delta=0.1)
+
+    def test_index_length_mismatch(self):
+        """A tensor ``index`` must hold one class index per image in the batch."""
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        model = DENSENET3D.to(device)
+        model.eval()
+        x = torch.rand((3, 1, 6, 6, 6), device=device)
+        with self.assertRaises(ValueError):
+            VanillaGrad(model)(x, index=torch.tensor([0, 1], device=device))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/visualize/test_vis_gradcam.py

@@ -19,7 +19,7 @@
 from parameterized import parameterized
 
 from monai.networks.nets import DenseNet, DenseNet121, SEResNet50
-from monai.visualize import GradCAM, GradCAMpp
+from monai.visualize import GradCAM, GradCAMpp, ModelWithHooks
 from tests.test_utils import assert_allclose, skip_if_quick
 
 
@@ -146,6 +146,39 @@ def __call__(self, x, adjoint_info):
 
     TESTS_ILL.append([cam])
 
+TESTS_BATCH: list[Any] = [[GradCAM], [GradCAMpp]]
+
+
+def _spatial_output_net() -> torch.nn.Module:
+    """A small network with spatial logits ``(B, 3, H, W)``; its layer ``"1"`` keeps the spatial size of the input."""
+    return torch.nn.Sequential(torch.nn.Conv2d(1, 4, 3, padding=1), torch.nn.ReLU(), torch.nn.Conv2d(4, 3, 1))
+
+
+class TestClassScore(unittest.TestCase):
+    def test_class_score(self):
+        """``class_score`` returns one score per batch element for integer and per-element class indices."""
+        hooks = ModelWithHooks(torch.nn.Linear(2, 4), target_layer_names=())
+        logits = torch.arange(12.0).reshape(3, 4)
+        assert_allclose(hooks.class_score(logits, 2), logits[:, 2])
+        assert_allclose(hooks.class_score(logits, torch.tensor(2)), logits[:, 2])
+        assert_allclose(hooks.class_score(logits, torch.tensor([0, 3, 1])), torch.tensor([0.0, 7.0, 9.0]))
+        # one class per batch element is broadcast over the trailing (spatial) dimensions of the logits
+        spatial = torch.arange(24.0).reshape(2, 3, 2, 2)
+        assert_allclose(hooks.class_score(spatial, torch.tensor([2, 0])), torch.stack([spatial[0, 2], spatial[1, 0]]))
+        with self.assertRaises(ValueError):
+            hooks.class_score(logits, torch.tensor([0, 1]))
+
+    def test_training_mode_restored_on_error(self):
+        """A class index of the wrong length raises ``ValueError`` and leaves the model in its previous mode."""
+        model = _spatial_output_net()
+        cam = GradCAM(nn_module=model, target_layers="1")
+        image = torch.rand((2, 1, 8, 8))
+        for training in (True, False):
+            model.train(training)
+            with self.assertRaises(ValueError):
+                cam(x=image, class_idx=torch.tensor([0, 1, 2]))
+            self.assertEqual(model.training, training)
+
 
 @skip_if_quick
 class TestGradientClassActivationMap(unittest.TestCase):
@@ -195,6 +228,49 @@ def test_shape(self, cam_class, input_data, expected_shape):
         result2 = cam(x=image, layer_idx=-1, class_idx=inferred, **kwargs)
         assert_allclose(result, result2)
 
+    @parameterized.expand(TESTS_BATCH)
+    def test_batch_matches_single_sample(self, cam_class):
+        """
+        The maps of a batch equal the maps computed one image at a time, for ``class_idx`` given as
+        ``None`` (per-image argmax), an integer (same class for every image) and a tensor (one class per image).
+        """
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        model = DenseNet(
+            spatial_dims=3, in_channels=1, out_channels=3, init_features=2, growth_rate=2, block_config=(6,)
+        )
+        model.to(device)
+        model.eval()
+        cam = cam_class(nn_module=model, target_layers="class_layers.relu")
+        image = torch.rand((3, 1, 6, 6, 6), device=device)
+        for class_idx in (None, 1, torch.tensor([0, 2, 1], device=device)):
+            batch = cam(x=image, class_idx=class_idx)
+            single = []
+            for i in range(image.shape[0]):
+                class_idx_i = int(class_idx[i]) if isinstance(class_idx, torch.Tensor) else class_idx
+                single.append(cam(x=image[i : i + 1], class_idx=class_idx_i))
+            assert_allclose(batch, torch.cat(single), rtol=1e-4, atol=1e-5)
+
+    @parameterized.expand(TESTS_BATCH)
+    def test_spatial_logits(self, cam_class):
+        """
+        Models with spatial logits ``(B, C, H, W)`` give one finite map per image, equal to the map of each image
+        on its own, for ``class_idx`` given as ``None`` (per-location argmax), an integer and a per-image tensor.
+        """
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        model = _spatial_output_net().to(device)
+        model.eval()
+        cam = cam_class(nn_module=model, target_layers="1")
+        image = torch.rand((2, 1, 8, 8), device=device)
+        for class_idx in (None, 1, torch.tensor([2, 0], device=device)):
+            batch = cam(x=image, class_idx=class_idx)
+            self.assertTupleEqual(batch.shape, (2, 1, 8, 8))
+            self.assertTrue(torch.isfinite(batch).all())
+            single = []
+            for i in range(image.shape[0]):
+                class_idx_i = int(class_idx[i]) if isinstance(class_idx, torch.Tensor) else class_idx
+                single.append(cam(x=image[i : i + 1], class_idx=class_idx_i))
+            assert_allclose(batch, torch.cat(single), rtol=1e-4, atol=1e-5)
+
     @parameterized.expand(TESTS_ILL)
     def test_ill(self, cam_class):
         model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3)