Pytorch uses hook to get feature map

For more detailed suggestions, please refer to Use hook Get feature map visualization

Define the following two hook function , The feature map can be obtained in a few steps

fmap_block = dict()  #  loading feature map
grad_block = dict()  #  Loading gradient 

def backward_hook(module, grad_in, grad_out):
    grad_block['grad_in'] = grad_in
    grad_block['grad_out'] = grad_out


def farward_hook(module, inp, outp):
    fmap_block['input'] = inp
    fmap_block['output'] = outp

import torch
from torch import nn

# Define the network 
class LeNet(nn.Module):  
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3,6,3,1,1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6,9,3,1,1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2,2)
        self.fc1 = nn.Linear(8*8*9, 120)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(120,10)

    def forward(self, x):
        out = self.pool1(self.relu1(self.conv1(x)))
        out = self.pool2(self.relu2(self.conv2(out)))
        out = out.view(out.shape[0], -1)
        out = self.relu3(self.fc1(out))
        out = self.fc2(out)

        return out

# Instantiate the network 
net = LeNet()


#  register hook
net.conv2.register_forward_hook(farward_hook)

 Define input 

input_img = torch.randn(1,3,32,32).requires_grad_()  
fmap_block = dict()  #  loading feature map
outs = net(input_img)

print(len(fmap_block))
print(fmap_block["output"].shape)