【Pytorch】使用Hook获取网络中间输出

在一些情况下（debug，绘制热力图）我们需要获取神经网络的某一层的输出。
如果在模型中直接增加一个return，将值返回是一种效率很低的办法。
在Pytorch中内置了Hook方法，可以帮助我们实现获取数据。

注册hook

首先需要在class中注册hook(class需要继承nn.Module)，注册方法如下:

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super(MLP,self).__init__()
        self.fc1 = nn.Linear(784,256)
        self.fc2 = nn.Linear(256,10)
        
        #* 用于在hook中承载数据
        self.gradients = list()
        self.conv_outputs = list()
        
        #* 注册前向hook
        self.fc1.register_forward_hook(self.forward_hook)
        #* 注册反向hook
        self.fc2.register_full_backward_hook(self.backward_hook)
        
    def forward(self,x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

定义hook函数

在注册hook，pytorch会在前向经过fc1，反向经过fc2时调用hook函数，上述的代码只注册了hook，但是没有定义hook函数，下面定义hook函数。

    def forward_hook(self, module, input, output):
        self.conv_outputs.append(output)

    def backward_hook(self, module, grad_in, grad_out):
        self.gradients.append(grad_out[0])

获取hook到数据

经过上述两步，网络在被训练时会自动记录每一次的fc1前向传播结果以及每一次fc2反向传播结果。
我们只需要通过model.gradients以及model.conv_outputs即可获取到这两个数组。

完整代码

注，此处代码只是示例，展示一下hook方法的全流程

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super(MLP,self).__init__()
        self.fc1 = nn.Linear(784,256)
        self.fc2 = nn.Linear(256,10)
        
        #* 用于在hook中承载数据
        self.gradients = list()
        self.conv_outputs = list()
        
        #* 注册前向hook
        self.fc1.register_forward_hook(self.forward_hook)
        #* 注册反向hook
        self.fc2.register_full_backward_hook(self.backward_hook)

    def forward_hook(self, module, input, output):
        self.conv_outputs.append(output)

    def backward_hook(self, module, grad_in, grad_out):
        self.gradients.append(grad_out[0])

    def forward(self,x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x
    

if __name__=='__main__':
    model=MLP()
    x = torch.randn(1,784,requires_grad=True)
    y = model(x)
    y.backward(torch.ones([1,10]))

    #* 获取数据
    gradients=model.gradients
    conv_outputs=model.conv_outputs

    print(gradients)
    print(conv_outputs)