更细节的建议参考使用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)