How to use GradCAM inside pytorch_lightning

[kapsner]

Hello,

I am trying to use monai.visualize.GradCAM during inference in pytorch lightning but constantly run into this error:

Traceback (most recent call last):
File "/home/user/.vscode-server/extensions/ms-python.python-2021.2.576481509/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py", line 416, in evaluate_expression
compiled = compile(_expression_to_evaluate(expression), '', 'eval')
File "", line 1
gradcam_img = self.gradcam(x=raw_image[None], class_idx=None)
^
SyntaxError: invalid syntax

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "", line 1, in
File "/home/user/miniconda/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 339, in call
acti_map = self.compute_map(x, class_idx=class_idx, retain_graph=retain_graph, layer_idx=layer_idx)
File "/home/user/miniconda/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 305, in compute_map
logits, acti, grad = self.net(x, class_idx=class_idx, retain_graph=retain_graph)
File "/home/user/miniconda/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 112, in call
score.sum().backward(retain_graph=retain_graph)
File "/home/user/miniconda/lib/python3.8/site-packages/torch/tensor.py", line 221, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/home/user/miniconda/lib/python3.8/site-packages/torch/autograd/init.py", line 130, in backward
Variable._execution_engine.run_backward(
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

The relevant code (using a densnet121 model) is as follows:

class BackboneModel(pl.LightningModule):
    # https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html
    def __init__(self, conf):
        super().__init__()

        self.model = monai.networks.nets.densenet.densenet121(
                spatial_dims=2,
                in_channels=1,
                out_channels=2
            )
        self.gradcam = monai.visualize.GradCAM(
            nn_module=self.model,
            target_layers="class_layers.relu"
        )

      def forward(self, x):
        return self.model(x)

      def train_step(self, batch, batch_idx, prefix):
        x = batch["image"]
        y = batch["target"]

        logits = self(x)
        return {"target": y, "logits": logits}
 
      def validation_step(self, batch, batch_idx, prefix):
        x = batch["image"]
        y = batch["target"]

        logits = self(x)
        return {"target": y, "logits": logits}    
      
      def test_step(self, batch, batch_idx, prefix):
        x = batch["image"]
        y = batch["target"]

        logits = self(x)
        
        raw_image = x[0]
        gradcam_img = self.gradcam(x=raw_image[None], class_idx=None)

Any idea how to properly incorporate GradCAM into this setup?

Note: x[0].requires_grad returns False

[rijobro]

I suspect that GradCAM is returning a numpy array, whereas a torch.Tensor is required for lightning. What happens if you modify this:

gradcam_img = self.gradcam(x=raw_image[None], class_idx=None)

to this:

gradcam_img = torch.Tensor(self.gradcam(x=raw_image[None], class_idx=None))

[kapsner]

Hi @rijobro,

thanks for your quick reply - I appreciate that very much :)

Unfortunatelly, the error stays the same:

torch.Tensor(self.gradcam(
x=raw_image[None], class_idx=None))
Traceback (most recent call last):
File "", line 1, in
File "/home/user/miniconda/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 339, in call
acti_map = self.compute_map(x, class_idx=class_idx, retain_graph=retain_graph, layer_idx=layer_idx)
File "/home/user/miniconda/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 305, in compute_map
logits, acti, grad = self.net(x, class_idx=class_idx, retain_graph=retain_graph)
File "/home/user/miniconda/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 112, in call
score.sum().backward(retain_graph=retain_graph)
File "/home/user/miniconda/lib/python3.8/site-packages/torch/tensor.py", line 221, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/home/user/miniconda/lib/python3.8/site-packages/torch/autograd/init.py", line 130, in backward
Variable._execution_engine.run_backward(
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

I wonder, if the issue could somehow be related to the model state when running GradCAM. As far as I understand it, pytorch_lightning ensures that within "test_step", the model is in "inference state" (aka. model.eval()).

[wyli]

I'm not familiar with pytorch lightning but it's fine if x[0].requires_grad returns False. please ensure for x in self.model.parameters(): print(x.requires_grad) have some True... could try to put this line gradcam_img = self.gradcam(x=raw_image[None], class_idx=None) in the training step?

[kapsner]

Would it make sense to run the GradCAM during training? I am currently predicting classes of test images and want to support interpretability using the GradCAM algorithm.

By the way: self.model.parameters(): print(x.requires_grad) returns True for all parameters just at the position in the code where gradcam is called.

for x in self.model.parameters(): print(x.requires_grad)
True
True
True
True
True
True
True
True
True
True
True
...

[kapsner]

I have created a minimal working example in a kaggle notebook using pytorch lightning's bug report model.

@wyli you were right: calling GradCAM inside train_step works. So I assume the error is related to the model state inside the test_step.

But does it really make sense to run a class activation map on training examples? My understanding is that this should more meaningful and possible from inside test_step...

[rijobro]

What happens if you change this:

cam_img = self.cam(x=raw_image[None], class_idx=None)

to this:

with torch.set_grad_enabled(True):
    cam_img = self.cam(x=raw_image[None], class_idx=None)

[rijobro]

I copied your kaggle notebook to google colab (so I could enable GPU), and the suggested change seems to work for me. Let me know if it doesn't for you.

[kapsner]

@rijobro your solution works, many thanks!!

Would torch.set_grad_enabled(True) have any unwanted side effects when using it inside test_step to get GradCAM working, e.g. changing/updating the model's weights?

[rijobro]

side effects when using it inside trains_step

Do you mean test_step?

No I don't think so. There would be no difference between calculating the gradcam in train_step or test_step, it's just that gradients are disabled by default in test_step.

[kapsner]

Sorry for the typo - yes, I meant test_step!

Thank you very much for your help!! Appreciate the quick responses very much!!

[MJGroeneveld]

[KumoLiu]

Hi @kapsner, you may need to modify this part manually.

MONAI/monai/visualize/class_activation_maps.py

Lines 134 to 135 in 0ffe55a

	logits = self.model(x, **kwargs)
	self.class_idx = logits.max(1)[-1] if class_idx is None else class_idx

Something like:

classification_output, bounding_box_output = self.model(x, **kwargs)
self.class_idx = classification_output.max(1)[-1] if class_idx is None else class_idx

Hope it can help you, thanks!