PyTorch之HOOK——获取神经网络特征和梯度的有效工具

本文首发于简书 西北小生_ 的博客：https://www.jianshu.com/u/898c7641f6ea，未经允许，禁止转载！

为了更深入地理解神经网络模型，有时候我们需要观察它训练得到的卷积核、特征图或者梯度等信息，这在CNN可视化研究中经常用到。其中，卷积核最易获取，将模型参数保存即可得到；特征图是中间变量，所对应的图像处理完即会被系统清除，否则将严重占用内存；梯度跟特征图类似，除了叶子结点外，其它中间变量的梯度都被会内存释放，因而不能直接获取。
最容易想到的获取方法就是改变模型结构，在forward的最后不但返回模型的预测输出，还返回所需要的特征图等信息。

如何在不改变模型结构的基础上获取特征图、梯度等信息呢？

Pytorch的hook编程可以在不改变网络结构的基础上有效获取、改变模型中间变量以及梯度等信息。
hook可以提取或改变Tensor的梯度，也可以获取nn.Module的输出和梯度（这里不能改变）。因此有3个hook函数用于实现以上功能：

Tensor.register_hook(hook_fn)，

nn.Module.register_forward_hook(hook_fn)，

nn.Module.register_backward_hook(hook_fn).

下面对其用法进行一一介绍。

1.Tensor.register_hook(hook_fn)

功能：注册一个反向传播hook函数，用于自动记录Tensor的梯度。
PyTorch对中间变量和非叶子节点的梯度运行完后会自动释放，以减缓内存占用。什么是中间变量？什么是非叶子节点？

Tensor计算

上图中，a，b，d就是叶子节点，c，e，o是非叶子节点，也是中间变量。

In [18]: a = torch.Tensor([1,2]).requires_grad_() 
    ...: b = torch.Tensor([3,4]).requires_grad_() 
    ...: d = torch.Tensor([2]).requires_grad_() 
    ...: c = a + b 
    ...: e = c * d 
    ...: o = e.sum()     

In [19]: o.backward()

In [20]: print(a.grad)
tensor([2., 2.])

In [21]: print(b.grad)
tensor([2., 2.])

In [22]: print(c.grad)
None

In [23]: print(d.grad)
tensor([10.])

In [24]: print(e.grad)
None

In [25]: print(o.grad)
None

可以从程序的输出中看到，a，b，d作为叶子节点，经过反向传播后梯度值仍然保留，而其它非叶子节点的梯度已经被自动释放了，要想得到它们的梯度值，就需要使用hook了。

我们首先自定义一个hook_fn函数，用于记录对Tensor梯度的操作，然后用Tensor.register_hook(hook_fn)对要获取梯度的非叶子结点的Tensor进行注册，然后重新反向传播一次：

In [44]: def hook_fn(grad):
    ...:     print(grad)
    ...:

In [45]: e.register_hook(hook_fn)
Out[45]: <torch.utils.hooks.RemovableHandle at 0x1d139cf0a88>

In [46]: o.backward()
tensor([1., 1.])

这时就自动输出了e的梯度。

自定义的hook_fn函数的函数名可以是任取的，它的参数是grad，表示Tensor的梯度。这个自定义函数主要是用于描述对Tensor梯度值的操作，上例中我们是对梯度直接进行输出，所以是print(grad)。我们也可以把梯度装在一个列表或字典里，甚至可以修改梯度，这样如果梯度很小的时候将其变大一点就可以防止梯度消失的问题了：

In [28]: a = torch.Tensor([1,2]).requires_grad_() 
    ...: b = torch.Tensor([3,4]).requires_grad_() 
    ...: d = torch.Tensor([2]).requires_grad_() 
    ...: c = a + b 
    ...: e = c * d 
    ...: o = e.sum()                                                            

In [29]: grad_list = []                                                         

In [30]: def hook(grad): 
    ...:     grad_list.append(grad)    # 将梯度装在列表里
    ...:     return 2 * grad    # 将梯度放大两倍
    ...:                                                                        

In [31]: c.register_hook(hook)                                                  
Out[31]: <torch.utils.hooks.RemovableHandle at 0x7f009b713208>

In [32]: o.backward()                                                           

In [33]: grad_list                                                              
Out[33]: [tensor([2., 2.])]

In [34]: a.grad                                                                 
Out[34]: tensor([4., 4.])

In [35]: b.grad                                                                 
Out[35]: tensor([4., 4.])

上例中，我们定义的hook函数执行了两个操作：一是将梯度装进列表grad_list中，二是把梯度放大两倍。从输出中我们可以看到，执行反向传播后，我们注册的非叶子节点c的梯度保存在了列表grad_list中，并且a和b的梯度都变为原来的两倍。这里需要注意的是，如果要将梯度值装在一个列表或字典里，那么首先要定义一个同名的全局变量的列表或字典，即使是局部变量，也要在自定义的hook函数外面。另一个需要注意的点就是如果要改变梯度值，hook函数要有返回值，返回改变后的梯度。

这里总结一下，如果要获取非叶子节点Tensor的梯度值，我们需要在反向传播前：
1）自定义一个hook函数，描述对梯度的操作，函数名自拟，参数只有grad，表示Tensor的梯度；
2）对要获取梯度的Tensor用方法Tensor.register_hook(hook)进行注册。
3）执行反向传播。

2.nn.Module.register_forward_hook(hook_fn)和nn.Module.register_backward_hook(hook_fn)

这两个的操作对象都是nn.Module类，如神经网络中的卷积层(nn.Conv2d)，全连接层(nn.Linear)，池化层(nn.MaxPool2d, nn.AvgPool2d)，激活层(nn.ReLU)或者nn.Sequential定义的小模块等，所以放在一起讲。

对于模型的中间模块，也可以视作中间节点(非叶子节点)，它的输出为特征图或激活值，反向传播的梯度值都会被系统自动释放，如果想要获取它们，就要用到hook功能。

有名字即可看出，register_forward_hook是获取前向传播的输出的，即特征图或激活值；register_backward_hook是获取反向传播的输出的，即梯度值。它们的用法和上面介绍的register_hook类似。我们先看一下hook_fn的定义：

对于register_forward_hook(hook_fn)，其hook_fn函数定义如下：

def forward_hook(module, input, output):
    operations

这里有3个参数，分别表示：模块，模块的输入，模块的输出。函数用于描述对这些参数的操作，一般我们都是为了获取特征图，即只描述对output的操作即可。

对于register_backward_hook(hook_fn)，其hook_fn函数定义如下：

def backward_hook(module, grad_in, grad_out):
    operations

这里也有3个参数，分别表示：模块，模块输入端的梯度，模块输出端的梯度。这里需要特别注意的是，此处的输入端和输出端，是前向传播时的输入端和输出端，也就是说，上面的output的梯度对应这里的grad_out。例如线性模块：o=W*x+b，其输入端为 W，x 和 b，输出端为 o。

如果模块有多个输入或者输出的话，grad_in和grad_out可以是 tuple 类型。对于线性模块：o=W*x+b ，它的输入端包括了W、x 和 b 三部分，因此 grad_input 就是一个包含三个元素的 tuple。

这里注意和 forward hook 的不同：

1. 在 forward hook 中，input 是 x，而不包括 W 和 b。
2. 返回 Tensor 或者 None，backward hook 函数不能直接改变它的输入变量，但是可以返回新的 grad_in，反向传播到它上一个模块。

此处的自定义的函数hook_fn也可以自拟名称，但如果两个hook函数同时使用的时候注意名称的区别，一般在函数名里添加对应的forward和backward就不易搞混了。

下面看一个具体用例：

#-*- utf-8 -*-

'''本程序用于验证hook编程获取卷积层的输出特征图和特征图的梯度'''

__author__ = 'puxitong from UESTC'

import torch
import torch.nn as nn
import numpy as np 
import torchvision.transforms as transforms


class Net(nn.Module):
    def __init__(self):
        super(Net, 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


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


loss_func = nn.CrossEntropyLoss()

# 生成一个假标签以便演示
label = torch.empty(1, dtype=torch.long).random_(3)

# 生成一副假图像以便演示
input_img = torch.randn(1,3,32,32).requires_grad_()  

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

net = Net()

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

outs = net(input_img)
loss = loss_func(outs, label)
loss.backward()

print('End.')

上面的程序中，我们先定义了一个简单的卷积神经网络模型，我们对第二层卷积模块进行hook注册，既获取它的输入输出，又获取输入输出的梯度，并将它们分别装在字典里。为了达到验证效果，我们随机生成一个假图像，它的尺寸和cifar-10数据集的图像尺寸一致，并且给这个假图像定义一个类别标签，用损失函数进行反向传播，模拟神经网络的训练过程。

在IPython中运行程序后，相应的特征图和梯度就会出现在两个列表fmap_block和grad_block中了。我们看一下它们的输入和输出的维度：

In [17]: len(fmap_block['input'])                                               
Out[17]: 1

In [18]: len(fmap_block['output'])                                              
Out[18]: 1

In [19]: len(grad_block['grad_in'])                                             
Out[19]: 3

In [20]: len(grad_block['grad_out'])                                            
Out[20]: 1

可以看出，第二层卷积模块的输入和输出都只有一个，即相应的特征图。而输入端的梯度值有3个，分别为权重的梯度，偏差的梯度，以及输入特征图的梯度。输出端的梯度值只有一个，即输出特征图的梯度。正如上面强调的，输入端即使有W, X和b三个，对于前项传播来说只有X是其输入，而对于反向传播来说，3个都是输入。输出端3项的梯度值排列的顺序是什么呢，我们来看一下3项梯度的具体维度：

In [21]: grad_block['grad_in'][0].shape                                         
Out[21]: torch.Size([1, 6, 16, 16])

In [22]: grad_block['grad_in'][1].shape                                         
Out[22]: torch.Size([9, 6, 3, 3])

In [23]: grad_block['grad_in'][2].shape                                         
Out[23]: torch.Size([9])

从输出端梯度的维度可以判断，第一个显然是特征图的梯度，第二个则是权重(卷积核/滤波器)的梯度，第三个是偏置的梯度。为了验证梯度和这些参数具有同样的维度，我们再来看看这三个值前向传播时的维度：

In [24]: fmap_block['input'][0].shape                                           
Out[24]: torch.Size([1, 6, 16, 16])

In [25]: net.conv2.weight.shape         
Out[25]: torch.Size([9, 6, 3, 3])

In [26]: net.conv2.bias.shape                                                   
Out[26]: torch.Size([9])

可以看到，我们的判断是正确的。

最后需要注意的一点是，如果需要获取输入图像的梯度，一定要将输入Tensor的requires_grad属性设为True。

原创不易，有用请点赞支持~