深度学习中的四种钩子

为了节省显存（内存），pytorch在计算过程中不保存中间变量，包括中间层的特征图和非叶子张量的梯度等。有时对网络进行分析时需要查看或修改这些中间变量，此时就需要注册一个钩子（hook）来导出需要的中间变量。网上介绍这个的有不少，但我看了一圈，多少都有不准确或不易懂的地方，我这里再总结一下，给出实际用法和注意点。
hook方法有四种:
torch.Tensor.register_hook()
torch.nn.Module.register_forward_hook()
torch.nn.Module.register_backward_hook()
torch.nn.Module.register_forward_pre_hook().

第一种torch.Tensor.register_hook()

import torch
def grad_hook(grad):
    grad *= 1.6
x = torch.tensor([1., 1., 1., 1.], requires_grad=True)
y = torch.pow(x, 2)
z = torch.sum(y)
h = x.register_hook(grad_hook)
z.backward()
print(x.grad)
h.remove()    # removes the hook

结果为：

tensor([3.2000, 3.2000, 3.2000, 3.2000])

import torch
def grad_hook(grad):
    grad *= 50
x = torch.tensor([2., 2., 2., 2.], requires_grad=True)
y = torch.pow(x, 2)
z = torch.mean(y)
h = x.register_hook(grad_hook)
z.backward()
print(x.grad)
h.remove()    # removes the hook
>>> tensor([50., 50., 50., 50.])

这个值是咋算的呢，其实是在反向传播的时候，获取x的梯度，上述代码不光获取了x的梯度，还将其乘以1.6，能够改变梯度的值。
注：

可以用remove()方法取消hook。注意remove()必须在backward()之后，因为只有在执行backward()语句时，pytorch才开始计算梯度，而在x.register_hook(grad_hook)时它仅仅是"注册"了一个grad的钩子，此时并没有计算，而执行remove就取消了这个钩子，然后再backward()时钩子就不起作用了。

第二种 torch.nn.Module.register_forward_hook(module, in, out)
用来导出指定子模块（可以是层、模块等nn.Module类型）的输入输出张量，但只可修改输出，常用来导出或修改卷积特征图。

inps, outs = [],[]
def layer_hook(module, inp, out):
    inps.append(inp[0].data.cpu().numpy())
    outs.append(out.data.cpu().numpy())

hook = net.layer1.register_forward_hook(layer_hook)
output = net(input)
hook.remove()

注意：（1）因为模块可以是多输入的，所以输入是tuple型的，需要先提取其中的Tensor再操作；输出是Tensor型的可直接用。
   （2）导出后不要放到显存上，除非你有A100。
   （3）只能修改输出out的值，不能修改输入inp的值（不能返回，本地修改也无效），修改时最好用return形式返回，如：

def layer_hook(self, module, inp, out):
    out = self.lam * out + (1 - self.lam) * out[self.indices]
    return out

这段代码用在manifold mixup中，用来对中间层特征进行混合来实现数据增强，其中self.lam是一个[0,1]概率值，self.indices是shuffle后的序号。

3, torch.nn.Module.register_forward_pre_hook(module, in)
用来导出或修改指定子模块的输入张量。

def pre_hook(module, inp):
    inp0 = inp[0]
    inp0 = inp0 * 2
    inp = tuple([inp0])
    return inp

hook = net.layer1.register_forward_pre_hook(pre_hook)
output = net(input)
hook.remove()

注意：（1）inp值是个tuple类型，所以需要先把其中的张量提取出来，再做其他操作，然后还要再转化为tuple返回。
（2）在执行output = net(input)时才会调用此句，remove()可放在调用后用来取消钩子。

4, torch.nn.Module.register_backward_hook(module, grad_in, grad_out)
  用来导出指定子模块的输入输出张量的梯度，但只可修改输入张量的梯度（即只能返回gin），输出张量梯度不可修改。

gouts = []
def backward_hook(module, gin, gout):
    print(len(gin),len(gout))
    gouts.append(gout[0].data.cpu().numpy())
    gin0,gin1,gin2 = gin
    gin1 = gin1*2
    gin2 = gin2*3
    gin = tuple([gin0,gin1,gin2])
    return gin

hook = net.layer1.register_backward_hook(backward_hook)
loss.backward()
hook.remove()

注意：
（1）其中的grad_in和grad_out都是tuple，必须要先解开，修改时执行操作后再重新放回tuple返回。
（2）这个钩子函数在backward()语句中被调用，所以remove()要放在backward()之后用来取消钩子。