为了更深入地理解神经网络模型，有时候我们需要观察它训练得到的卷积核、特征图或者梯度等信息，这在CNN可视化研究中经常用到。其中，卷积核最易获取，将模型参数保存即可得到；特征图是中间变量，所对应的图像处理完即会被系统清除，否则将严重占用内存；梯度跟特征图类似，除了叶子结点外，其它中间变量的梯度都被会内存释放，因而不能直接获取。
最容易想到的获取方法就是改变模型结构，在forward的最后不但返回模型的预测输出，还返回所需要的特征图等信息。
那么如何在不改变模型结构的基础上(比如对于预训练模型或者已训练好的模型)获取特征图、梯度等信息呢？
Pytorch的hook编程可以在不改变网络结构的基础上有效获取、改变模型中间变量以及梯度等信息。
hook可以提取或改变Tensor的梯度，也可以获取 nn.Module 的输出和梯度（这里不能改变）。
因此有3个hook函数用于实现以上功能：
(1) Tensor对象:
Tensor.register_hook(hook)
(2) Module对象:
nn.Module.register_forward_hook(hook) ;
nn.Module.register_backward_hook(hook).
下面对其用法进行一一介绍：
Tensor对象
官方对于Tensor对象的介绍：https://pytorch.org/docs/stable/tensors.html
有如下的 register_hook(hook) 方法，为Tensor注册一个 backward hook，用来获取变量的梯度。
hook必须遵循如下的格式：hook(grad) -> Tensor or None，其中grad为获取的梯度。
功能：注册一个反向传播hook函数，用于自动记录Tensor的梯度。
PyTorch对中间变量和非叶子节点的梯度运行完后会自动释放，以减缓内存占用。什么是中间变量？什么是非叶子节点？
上图中，a，b，d就是叶子节点，c，e，o是非叶子节点，也是中间变量。
In [1]: 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 [2]: o.backward()
In [3]: print(a.grad)
Out[3]: tensor([2., 2.])
In [4]: print(b.grad)
Out[4]: tensor([2., 2.])
In [5]: print(c.grad)
Out[5]: None
In [6]: print(d.grad)
Out[6]: tensor([10.])
In [7]: print(e.grad)
Out[7]: None
In [8]: print(o.grad)
Out[8]: None
可以从程序的输出中看到，a，b，d作为叶子节点，经过反向传播后梯度值仍然保留，而其它非叶子节点的梯度已经被自动释放了，要想得到它们的梯度值，就需要使用hook了。
我们首先自定义一个hook函数，用于记录对Tensor梯度的操作，然后用 Tensor.register_hook(hook) 对要获取梯度的非叶子结点的Tensor进行注册，然后重新反向传播一次：
In [9]: def hook(grad):
...: print(grad)
...:
In [10]: e.register_hook(hook)
Out[10]: <torch.utils.hooks.RemovableHandle at 0x1d139cf0a88>
In [11]: o.backward()
In [12]: print(e.grad)
Out[12]: tensor([1., 1.])
这时就自动输出了e的梯度。
自定义的hook函数的函数名可以是任取的，它的参数是grad，表示Tensor的梯度。这个自定义函数主要是用于描述对Tensor梯度值的操作，上例中我们是对梯度直接进行输出，所以是print(grad)。我们也可以把梯度装在一个列表或字典里，甚至可以修改梯度，这样如果梯度很小的时候将其变大一点就可以防止梯度消失的问题了：
In [13]: 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 [14]: grad_list = []
In [15]: def hook(grad):
...: grad_list.append(grad) # 将梯度装在列表里
...: return 2 * grad # 将梯度放大两倍
...:
In [16]: c.register_hook(hook)
Out[16]: <torch.utils.hooks.RemovableHandle at 0x7f009b713208>
In [17]: o.backward()
In [18]: print(grad_list)
Out[18]: [tensor([2., 2.])]
In [19]: print(a.grad)
Out[19]: tensor([4., 4.])
In [20]: print(b.grad)
Out[20]: 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）执行反向传播。
Module对象
有 register_forward_hook(hook) 和 register_backward_hook(hook) 两种方法，分别对应前向传播和反向传播的hook函数。
这两个的操作对象都是nn.Module类，如神经网络中的卷积层(nn.Conv2d)，全连接层(nn.Linear)，池化层(nn.MaxPool2d, nn.AvgPool2d)，激活层(nn.ReLU)或者nn.Sequential定义的小模块等。
对于模型的中间模块，也可以视作中间节点(非叶子节点)，它的输出为特征图或激活值，反向传播的梯度值都会被系统自动释放，如果想要获取它们，就要用到hook功能。
有名字即可看出， register_forward_hook 是获取前向传播的输出的，即特征图或激活值； register_backward_hook 是获取反向传播的输出的，即梯度值。它们的用法和上面介绍的 register_hook 类似。hook 函数在使用后应及时删除，以避免每次都运行钩子增加运行负载。
(1) 对于 register_forward_hook(hook)，其hook函数定义如下：
# 这里有3个参数，分别表示：模块，模块的输入，模块的输出。
# 函数用于描述对这些参数的操作，一般我们都是为了获取特征图，即只描述对output的操作即可。
def forward_hook(module, input, output):
operations
hook可以修改input和output，但是不会影响forward的结果。最常用的场景是需要提取模型的某一层（不是最后一层）的输出特征，但又不希望修改其原有的模型定义文件，这时就可以利用forward_hook函数。
import torch
import torch.nn as nn
import torch.nn.functional as F
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
fmap_block = dict() # 装feature map
def forward_hook(module, input, output):
fmap_block['input'] = input
fmap_block['output'] = output
net = Net()
net.load_state_dict(torch.load('checkpoint.pth')) # 载入已训练好的模型
x = torch.randn(1, 3, 32, 32).requires_grad_() # 随机生成一副图像作为输入
handle = net.conv2.register_forward_hook(hook) # 注册hook
y = net(x)
# 展示输入图像和特定中间层的特征
plt.subplot(121)
plt.imshow(fmap_block['input'][0][0,0,:,:].cpu().detach().numpy())
plt.subplot(122)
plt.imshow(fmap_block['output'][0][0,0,:,:].cpu().detach().numpy())
plt.show()
print((fmap_block['input'][0].shape))
print((fmap_block['output'][0].shape))
handle.remove()
(2) 对于 register_backward_hook(hook)，其hook函数定义如下：
# 这里有3个参数，分别表示：模块，模块输入端的梯度，模块输出端的梯度。
def backward_hook(module, grad_in, grad_out):
operations
这里需要特别注意的是，此处的输入端和输出端，是前向传播时的输入端和输出端，也就是说，上面的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 的不同：
- 在 forward hook 中，input 是 x，而不包括 W 和 b。
- 返回Tensor 或者 None，backward hook 函数不能直接改变它的输入变量，但是可以返回新的
grad_in，反向传播到它上一个模块。
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
fmap_block = dict() # 装feature map
grad_block = dict() # 装梯度
def forward_hook(module, input, output):
fmap_block['input'] = input
fmap_block['output'] = output
def backward_hook(module, grad_in, grad_out):
grad_block['grad_in'] = grad_in
grad_block['grad_out'] = grad_out
loss_func = nn.CrossEntropyLoss()
label = torch.empty(1, dtype=torch.long).random_(3) # 生成一个假标签
input_img = torch.randn(1,3,32,32).requires_grad_() # 生成一副假图像作为输入
net = Net()
# 注册hook
handle_forward = net.conv2.register_forward_hook(farward_hook)
handle_backward = net.conv2.register_backward_hook(backward_hook)
outs = net(input_img)
loss = loss_func(outs, label)
loss.backward()
handle_forward.remove()
handle_backward.remove()
print('End.')
上面的程序中，我们先定义了一个简单的卷积神经网络模型，我们对第二层卷积模块进行hook注册，既获取它的输入输出，又获取输入输出的梯度，并将它们分别装在字典里。为了达到验证效果，我们随机生成一个假图像，它的尺寸和cifar-10数据集的图像尺寸一致，并且给这个假图像定义一个类别标签，用损失函数进行反向传播，模拟神经网络的训练过程。
运行程序后，相应的特征图和梯度就会出现在两个列表fmap_block 和 grad_block 中了。我们看一下它们的输入和输出的维度：
In [21]: print(len(fmap_block['input']))
Out[21]: 1
In [22]: print(len(fmap_block['output']))
Out[22]: 1
In [23]: print(len(grad_block['grad_in']))
Out[23]: 3
In [24]: print(len(grad_block['grad_out']))
Out[24]: 1
可以看出，第二层卷积模块的输入和输出都只有一个，即相应的特征图。而输入端的梯度值有3个，分别为权重的梯度，偏差的梯度，以及输入特征图的梯度。输出端的梯度值只有一个，即输出特征图的梯度。正如上面强调的，输入端即使有W, X和b三个，对于前项传播来说只有X是其输入，而对于反向传播来说，3个都是输入。输出端3项的梯度值排列的顺序是什么呢，我们来看一下3项梯度的具体维度：
In [25]: print(grad_block['grad_in'][0].shape)
Out[25]: torch.Size([1, 6, 16, 16])
In [26]: print(grad_block['grad_in'][1].shape)
Out[26]: torch.Size([9, 6, 3, 3])
In [27]: print(grad_block['grad_in'][2].shape)
Out[27]: torch.Size([9])
从输出端梯度的维度可以判断，第一个显然是特征图的梯度，第二个则是权重(卷积核/滤波器)的梯度，第三个是偏置的梯度。为了验证梯度和这些参数具有同样的维度，我们再来看看这三个值前向传播时的维度：
In [28]: print(fmap_block['input'][0].shape)
Out[28]: torch.Size([1, 6, 16, 16])
In [29]: print(net.conv2.weight.shape)
Out[29]: torch.Size([9, 6, 3, 3])
In [30]: print(net.conv2.bias.shape)
Out[30]: torch.Size([9])
最后需要注意的一点是，如果需要获取输入图像的梯度，一定要将输入Tensor的requires_grad属性设为True。