Generating the class activation maps

I have trained a CNN model with softmax classification layer. Now, I want to get the class activation maps (CAM) from these trained model on some test-samples. I have found a code that does this using Keras, but cannot do the same thing in PyTorch.

In this Keras code, they compute the gradients of the predicted output with respect to the last convolutional layer. So, I wonder how I should compute these gradients in PyTorch.

Here is the full Keras code:

    preds = model.predict(image)
    class_idx = np.argmax(preds[0])
    class_output = model.output[:, class_idx]
    last_conv_layer = model.get_layer("mixed10")

    grads = K.gradients(class_output, last_conv_layer.output)[0]
    pooled_grads = K.mean(grads, axis=(0, 1, 2))
    iterate = K.function([model.input], [pooled_grads, last_conv_layer.output[0]])
    pooled_grads_value, conv_layer_output_value = iterate([image])

    for i in range(2048):
        conv_layer_output_value[:, :, i] *= pooled_grads_value[i]

    heatmap = np.mean(conv_layer_output_value, axis = -1)
    heatmap = np.maximum(heatmap, 0)
    heatmap /= np.max(heatmap)

Any help is appreciate!

How about defining a hook? Would that solve your problem? E.g., something like

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
 
        self.gradients = []

        ## your layers
        self.some_conv_layer.register_backward_hook(self.save_gradients)
        
    def forward(self, x):
        ## your forward pass
        return F.softmax(logits)

    def save_gradients(self, module, grad_input, grad_output):
        self.gradients.append(grad_output)

model = MyModel()
output = model(x)
output.backward()

print(model.gradients)
3 Likes

Thank you very much! I think that solved the problem, and with that I should be able to get the activation maps.

Hi everyone , Is there a way to achieve the same using c++ frontend ?
Thank you.