Pytorch之Hook

Hook是Pytorch中一个十分有用的特性。利用它，我们可以不必改变网络输入输出的结构，方便地获取、改变网络中间层变量的值和梯度。这个功能被广泛用于可视化神经网络中间层的feature、gradient，从而诊断神经网络中可能出现的问题，分析网络有效性。

Pytorch中Hook的用法：

1. Hook for Tensors：针对Tensor的hook。
2. Hook for Modules：针对例如nn.Conv2d  nn.Linear等网络模块的hook。
3. Guided Backpropagation：利用hook实现神经网络可视化。

Hook for Tensors

在Pytorch的计算图中，只有叶子结点的变量会保留梯度，而所有中间变量的梯度只被用于反向传播，一旦完成反向传播，中间变量的梯度就会自动释放，从而节约内存。如果想要修改或者输出中间变量的梯度时，可以使用retain_grad()的方案，但是其保留梯度会增加内存占用，并不是个好办法。对此的一种替代方案就是用hook保存中间变量的梯度。

# z为中间变量
z.register_hook(lambda x: 2*x) # z的梯度扩大两倍
z.register_hook(lambda x: print(x)) # 输出z的梯度

Hook for Modules

网络模块module不像上一节的Tensor，拥有显示的变量名可以直接访问，而是被封装在神经网络中间。我们通常只能获得网络整体的输入和输出，对于夹在网络中间的模块，我们不但很难得知它输入/输出的梯度，甚至连它输入输出的数值都无法获得。除非设计网络时，在forward函数的返回值中包含中间module的输出，或者用很麻烦的办法，把网络按照module的名称拆分再组合，让中间层提供的feature暴露出来。

为了解决这个麻烦，Pytorch设计了两种hook：register_forward_hook和register_backward_hook，分别用来获取正/反向传播时，中间层模块输入和输出的feature/gradient，大大降低了获取模型内部信息流的难度。