Pytorch中hook如何直接调取网络中间层的输入输出数据（不理论分析，直接实践上手教程）

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

前言

本文记录了如何在pytorch中使用hook去调取网络中的任意一module的input和output数据
自己之前也摸索了很久很久，这里分享出来希望能帮助到大家！

---

一、hook是什么，以及怎么用？

在Pyorch中，可以利用Hook获取、改变网络中间某一层变量的值和梯度，从而便捷地分析网络，而不用专门改变网络结构，十分好用，这里只介绍如何用hook获取中间层变量值。

一共分为三步：

• 新建一个获取数据的函数（格式较为固定）
• 定位到自己要提取的网络中间层的位置
• 推理一遍网络，调用储存数据的变量以得到中间层数据

二、使用步骤

1.网络示例

网络代码如下（示例）：

import torch
import torch.nn as nn

#initial the network
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, 3, 1, 1),
            nn.ReLU(),
            nn.AvgPool2d(2, 2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(32 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU()
        )
        self.out = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        output = self.out(x)
        return output
        
net = ConvNet()

2.第一步：定义获取数据的函数

可根据自己需求更改函数内部代码，但最好函数名什么的不要变
代码如下（示例）：

input_data = {}
output_data = {}
def get_activation(name):
    def hook(model,input,output):
        input_data[name] = input[0].detach() # input type is tulple, only has one element, which is the tensor
        output_data[name] = output.detach()  # output type is tensor
    return hook

3.第二步：定位网络中间层位置

这一步是告诉hook哪些层的数据需要被调用
首先可以打印自己网络的结构

print(net)
#========the output is following ===============
ConvNet(
  (conv1): Sequential(
    (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): AvgPool2d(kernel_size=2, stride=2, padding=0)
  )
  (conv2): Sequential(
    (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (fc): Sequential(
    (0): Linear(in_features=1568, out_features=128, bias=True)
    (1): ReLU()
    (2): Linear(in_features=128, out_features=64, bias=True)
    (3): ReLU()
  )
  (out): Linear(in_features=64, out_features=10, bias=True)
)

比如说希望调用conv2中第0位的Conv2d这一层的数据以及fc中第1位的ReLU的数据
就可以在层位置后直接加上 .register_forward_hook(get_activation(‘xxx’))，其中xxx是方便后续将存储数据调用出来~

# hook the specific layers input and output
net.conv2[0].register_forward_hook(get_activation('conv2'))
net.fc[1].register_forward_hook(get_activation('ReLU'))

随后开始网络推理，给网络传入相应数据跑通一遍，这里只做演示，具体项目中还是得用相应数据集的数据~

x = torch.randn(1, 1, 28, 28)
y = net(x)
print('the data through conv2 is')
print(input_data['conv2']) # get the conv_2's input data
print(output_data['conv2']) # get the conv_2's input data
print('the data through ReLU is')
print(input_data['ReLU']) 
print(output_data['ReLU'])
#======output is following =================
the data through conv2 is
tensor([[[[0.0000e+00, 2.9236e-01, 2.0995e-01,  ..., 2.6224e-02,
           1.8165e-01, 1.7987e-01],
          [2.6442e-01, 2.0532e-01, 3.4805e-01,  ..., 3.5666e-02,
           0.0000e+00, 5.2180e-02],
          [1.8741e-01, 3.3035e-01, 1.8489e-01,  ..., 3.6469e-01,
           2.3763e-02, 1.6026e-01],
          ...,
          [4.0376e-01, 4.2911e-01, 2.8944e-01,  ..., 2.3436e-01,
           3.9946e-02, 2.6209e-01],
		............
		............
		............

由此就完成了hook对网络中特定层的数据调用且打印~

整体代码

import torch
import torch.nn as nn

#initial the network
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, 3, 1, 1),
            nn.ReLU(),
            nn.AvgPool2d(2, 2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(32 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU()
        )
        self.out = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        output = self.out(x)
        return output

# define the hook function
input_data = {}
output_data = {}
def get_activation(name):
    def hook(model,input,output):
        input_data[name] = input[0].detach() # input type is tulple, only has one element, which is the tensor
        output_data[name] = output.detach()  # output type is tensor
    return hook
net = ConvNet()

# hook the specific layers' input and output
net.conv2[0].register_forward_hook(get_activation('conv2'))
net.fc[1].register_forward_hook(get_activation('ReLU'))

x = torch.randn(1, 1, 28, 28)
y = net(x)
print('the data through conv2 is')
print(input_data['conv2']) # get the conv_2's input data
print(output_data['conv2']) # get the conv_2's input data
print('the data through ReLU is')
print(input_data['ReLU']) 
print(output_data['ReLU'])

---

总结

这篇总结了hook调用网络特定层的数据的步骤及用法，应该算是很精简易懂了，希望能帮助到大家~
若有疑问可以邮件博主：fuyz@stu.pku.edu.cn
冲！