pytorch使用hook获得特征图

更细节的建议参考使用hook获得特征图可视化

定义如下两个hook函数，简单几步就能获得特征图了

fmap_block = dict()  # 装feature map
grad_block = dict()  # 装梯度

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

#定义网路
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

#实例化网络
net = LeNet()


# 注册hook
net.conv2.register_forward_hook(farward_hook)

定义输入

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

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