Pytorch中的register_hook(梯度操作)

对于高阶调参师而言，对神经网络梯度级别的操作的不可避免的。有时候，咱们需要把某一层的梯度拿出来分析，辅助特征图可视化（如GradCAM）；再比如，hook还可以做优化器设计的实验。

hook，在中文里就是“钩子”的意思。Pytorch默认在反向传播过程中，不保留中间层的梯度，以达到减少内存的目的。如果需要对某一层的梯度特别感兴趣，咱们可以用钩子把它勾住，再等反向传播的时候，就可以访问到这一层的梯度。甚至可以对这层梯度进行修改，进而影响浅层的梯度传播。

我们先看官方给出的示例：

import torch
v = torch.tensor([0., 0., 0.], requires_grad=True)
h = v.register_hook(lambda grad: grad * 2)  # double the gradient
v.backward(torch.tensor([1., 2., 3.]))
print(v.grad)
h.remove()

这个栗子主要想表达register_hook可以用来修改梯度，register_hook的作用在v.backward()那里开始执行（众所周知，只有反向传播的时候才有梯度）。如第三行所表达的，在张量v这里对梯度放大2倍。
lambda其实就是一个无命名函数的表达，不了解的可以戳《python中的lambda关键字》。
通常，使用register_hook，一般形式是

tensor.register_hook(func)

这个func就是函数名，表示你想对这层的张量进行的操作。如果你只是想读取这一层的张量，那么你可以在func中用一个变量去缓存张量即可，比如：

cache = []
def func(grad):
    cache.append(grad)
    return grad

咱们可以在cache中获取到梯度了。因为即使是hook住的张量梯度，也会在运行结束后立刻释放掉，所以想要保存梯度，一定要在func函数里进行。

---

在神经网络中，咱们可以通过层的命名来获取该层的输出张量。比如对于检测器来讲，一般是backbone+head形式，假设咱们想获取backbone的梯度，可以考虑使用register_hook勾住backbone的输出张量，然后通过类似的方法保存住。

参考：https://pytorch.org/docs/stable/generated/torch.Tensor.register_hook.html
https://zhuanlan.zhihu.com/p/267130090