pytorch hook机制

为了更深入地理解神经网络模型，有时候我们需要观察它训练得到的卷积核、特征图或者梯度等信息，这在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对中间变量和非叶子节点的梯度运行完后会自动释放，以减缓内存占用。

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.])

参考：

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

model.apply(fn)或net.apply(fn) - 码农教程

pytorch的hook机制之register_forward_hook - 知乎

register_forward_hook_夜信_的博客-CSDN博客_register_forward_hook