0

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])
tanhz
  • 3
  • 1

1 Answers1

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

CuCaRot
  • 1,208
  • 7
  • 23