The activation map was introduced by Bolei Zhou et al. in “Learning Deep Features for Discriminative Localization”. It gives us some insight into why CNN made the predictions it did.

The activation map uses the output of the last convolutional layer (just before the average pooling layer) together with the predictions to give us a heatmap visualization of why the model made its decision. This is a useful tool for interpretation.

In this tutorial, we’ll learn “how to extract intermediate layer output from a convolutional neural network model using PyTorch hook”. I’ve created a small convolutional neural network for MNIST digit classification.

import torch
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return x
model = Net()
print(model)
PyTorch mnist model

We need a way to get access to the activations inside the mode. In PyTorch, this can be done with a hook. Hooks are PyTorch callbacks. However, rather than allowing you to inject code into the training loop, hooks allow you to inject code into the forward and backward calculations themselves. 

We can attach a hook to any layer of the model, and it will be executed when we compute the outputs (forward hook) or during backpropagation (backward hook).

A forward hook is a function that takes three things—a module, its input, and its output—and it can perform any behavior you want.

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

The name argument is used to store the activation in the dict.We put our hook function in a class so it has a state that we can access later, and just store a copy of the output. We can then instantiate a Hook and attach it to the layer we want, which is the first linear layer of the CNN body: 

hook=model.fc1.register_forward_hook(get_activation('fc1'))

Until now I have created a small code snippet using a forward hook to store one activation from fc1 and registered the activation with get_activation(‘fc2′). Now we can grab a batch and feed it through our model:

train_features, train_labels = next(iter(loaders['test']))
output = model(train_features)
print(activation['fc1'].shape)
print(activation['fc1'])

This would return the output of the registered module, so you would get the output of fc1 linear layer.

PyTorch Intermediate Layer output

If you would like to get the output of the F.relu, you could create an nn.ReLU() module and register a forward hook to this particular module or alternatively you could register a forward hook to the next module and store the input instead.

Once you’re done with your hook, you should remove it as otherwise it might leak some memory:

hook.remove()

Related Post