From 6df34e5d28aaf621a32aef0b5b8466d8b4b36edf Mon Sep 17 00:00:00 2001 From: Sonica-B Date: Sat, 24 Oct 2026 01:13:12 -0700 Subject: [PATCH 1/2] Support batched inputs in gradient-based saliency maps VanillaGrad.get_grad rejected inputs with more than one image because ModelWithHooks.class_score selected the class with logits[:, class_idx], which for one class index per image returns a (B, B) slice and mixes classes across the batch. class_score now accepts an int (same class for every image) or one index per image and returns one score per image (indexing one logit per image; torch.gather when a class is given per location), raising ValueError for an index of the wrong length. The batch check is removed, so VanillaGrad, SmoothGrad, GuidedBackpropGrad and GuidedBackpropSmoothGrad accept (B, C, ...) inputs; SmoothGrad scales its noise with the intensity range of each image. GradCAMpp reshapes the class score to (B, 1, ..., 1) before multiplying it with the gradient; on dev that product only broadcast when the batch size equalled the last feature-map dimension. Results for batch size 1 are unchanged. Tests compare batched maps with per-image maps for VanillaGrad, GuidedBackpropGrad, GradCAM and GradCAMpp with the index given as None, an int and a per-image tensor, and unit-test class_score directly. Fixes #7334 Signed-off-by: Sonica-B --- monai/visualize/class_activation_maps.py | 37 ++++++++++++++++---- monai/visualize/gradient_based.py | 16 ++++++--- tests/integration/test_vis_gradbased.py | 44 ++++++++++++++++++++++++ tests/visualize/test_vis_gradcam.py | 41 +++++++++++++++++++++- 4 files changed, 126 insertions(+), 12 deletions(-) diff --git a/monai/visualize/class_activation_maps.py b/monai/visualize/class_activation_maps.py index 39f26d0fbd4..c2ddef59d93 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,19 +124,42 @@ 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 + 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, cast(int, self.class_idx)) + 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: @@ -406,7 +429,9 @@ 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 + score = cast(torch.Tensor, self.nn_module.score).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..a5c74c1a5c2 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): @@ -38,14 +39,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 +71,43 @@ 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_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..e1493231401 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,23 @@ def __call__(self, x, adjoint_info): TESTS_ILL.append([cam]) +TESTS_BATCH: list[Any] = [[GradCAM], [GradCAMpp]] + + +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])) + @skip_if_quick class TestGradientClassActivationMap(unittest.TestCase): @@ -195,6 +212,28 @@ 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_ILL) def test_ill(self, cam_class): model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3) From 2f624a1e4214133f1f65b9f783e15ca0b76cae4a Mon Sep 17 00:00:00 2001 From: Sonica-B Date: Sun, 25 Oct 2026 03:08:54 -0700 Subject: [PATCH 2/2] Restore model mode on errors and reduce spatial GradCAM++ scores ModelWithHooks.__call__ restores the previous training mode in a finally block, so the ValueError raised for a class index of the wrong length, or any other error in the forward or backward pass, no longer leaves a training model in evaluation mode. For spatial logits (B, C, H, W) the class score holds one value per location, and the (B, 1, ..., 1) reshape added to GradCAMpp raised a RuntimeError. GradCAMpp now averages the score of each image first. 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 would overflow. On dev such models worked only for batch size 1 with an integer class index and a target layer with the spatial size of the output, where the score map was multiplied elementwise with the gradient; that case now uses the per-image score as well. Tests: a wrong-length index leaves the model in training or evaluation mode as it was; GradCAM and GradCAMpp on a network with spatial logits give finite maps equal to the per-image maps for class_idx None, an int and a per-image tensor; SmoothGrad with positive noise scales the noise of each image with its own intensity range, a check that fails with one range for the whole batch. Signed-off-by: Sonica-B --- monai/visualize/class_activation_maps.py | 44 +++++++++++++----------- tests/integration/test_vis_gradbased.py | 25 ++++++++++++++ tests/visualize/test_vis_gradcam.py | 37 ++++++++++++++++++++ 3 files changed, 86 insertions(+), 20 deletions(-) diff --git a/monai/visualize/class_activation_maps.py b/monai/visualize/class_activation_maps.py index c2ddef59d93..77706d269fb 100644 --- a/monai/visualize/class_activation_maps.py +++ b/monai/visualize/class_activation_maps.py @@ -152,24 +152,26 @@ def class_score(self, logits: torch.Tensor, class_idx: int | torch.Tensor) -> to def __call__(self, x, class_idx=None, retain_graph=False, **kwargs): train = self.model.training self.model.eval() - 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) - 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): @@ -429,8 +431,10 @@ 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) - # one class score per batch element, broadcast over the channel and spatial dimensions of its gradient - score = cast(torch.Tensor, self.nn_module.score).reshape(b, *([1] * (grad.ndim - 1))) + # 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) diff --git a/tests/integration/test_vis_gradbased.py b/tests/integration/test_vis_gradbased.py index a5c74c1a5c2..9693bc2bdb9 100644 --- a/tests/integration/test_vis_gradbased.py +++ b/tests/integration/test_vis_gradbased.py @@ -29,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) @@ -99,6 +107,23 @@ def test_smoothgrad_batch_without_noise(self): 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" diff --git a/tests/visualize/test_vis_gradcam.py b/tests/visualize/test_vis_gradcam.py index e1493231401..fc41db2298d 100644 --- a/tests/visualize/test_vis_gradcam.py +++ b/tests/visualize/test_vis_gradcam.py @@ -149,6 +149,11 @@ def __call__(self, x, adjoint_info): 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.""" @@ -163,6 +168,17 @@ def test_class_score(self): 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): @@ -234,6 +250,27 @@ def test_batch_matches_single_sample(self, cam_class): 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)