diff --git a/monai/visualize/class_activation_maps.py b/monai/visualize/class_activation_maps.py index 39f26d0fbd4..77706d269fb 100644 --- a/monai/visualize/class_activation_maps.py +++ b/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) diff --git a/monai/visualize/gradient_based.py b/monai/visualize/gradient_based.py index a8d81eb5790..1daef8cc3ad 100644 --- a/monai/visualize/gradient_based.py +++ b/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() diff --git a/tests/integration/test_vis_gradbased.py b/tests/integration/test_vis_gradbased.py index e9db0af2404..9693bc2bdb9 100644 --- a/tests/integration/test_vis_gradbased.py +++ b/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() diff --git a/tests/visualize/test_vis_gradcam.py b/tests/visualize/test_vis_gradcam.py index 274510ba30f..fc41db2298d 100644 --- a/tests/visualize/test_vis_gradcam.py +++ b/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)