How to get the gradients from loss to target layer output for Grad-CAM computation with pytorch

Question

I want to implement the code to get Grad-CAM map with pytorch(1.10.0). Most of implementation specify the target class to extract the gradients (This is a natural approach). But instead of this, I want to use gradients from loss to target layer output.

The following code is an example of a keras implementation of what I want to do. How can I implement it with pytorch?

loss = K.mean(K.square(y_pred - y_true), axis=-1)
output_conv = model.get_layer(layer_name).output
gradients = K.gradients(loss, output_conv)[0]

output, grads = K.function([model.input], [output_conv, gradients])([input_image])

Answer 1

You can register the hook for both the forward and backward of a model while inferencing. I wrote an easy example with ResNet50:

import torch
from torchvision import models

activation = {}
def get_forward(name):
    def hook(model, input, output):
        activation[name] = output.detach()
    return hook

def get_backward(name):
    def hook(model, input, output):
        activation[name] = output[0].detach()
    return hook

model = models.resnet50(pretrained=False)
model.layer4[2].register_forward_hook(get_forward('forward'))
model.layer4[2].register_backward_hook(get_backward('backward'))

x = torch.randn(10, 3, 224, 224)
out = model(x)
loss = out.mean()
loss.backward()

output, grads = activation['forward'], activation['backward']
print(output.shape)
print(grads.shape)

>>> torch.Size([10, 2048, 7, 7])
>>> torch.Size([10, 2048, 7, 7])

The function should work well with CNNs, be careful with the Linear layer where the hook doesn't work as expected because of this bug which is simplified by this answer