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

你可能感兴趣的