_register_hook

1 引言

由于在CAM与Grad-cam中存在提取模型中间层信息的步骤，需要使用hook机制，总结一下：

hook 的类型：

• tensor级别
• module级别

2 tensor级别的hook

register_hook的作用：即对x求导时，对x的导数进行操作，并且register_hook的参数只能以函数的形式传过去。

下面举个栗子：
#正常求导情况

实验一：tensor.register_hook()可以改变反向传播梯度值，对导数进行2倍操作
v = torch.randn((1, 3), dtype=torch.float32, requires_grad=True)
z = v.sum()
z.backward()
print(v.grad)
实验一输出：
tensor([[1., 1., 1.]])

实验一对照组
v = torch.randn((1, 3), dtype=torch.float32, requires_grad=True)
z = v.sum()
v.register_hook(lambda grad: grad*2)
z.backward()
print(v.grad)
tensor([[2., 2., 2.]]) ##实验一对照组 新的改变

实验二：tensor级别的register_hook函数，影响了整个模型的grad
x=Variable(torch.FloatTensor([[1, 2, 3]]),requires_grad=True)
mod=nn.Linear(3, 1, bias=False)

out=mod(x)
out.register_hook(lambda grad: 0.1*grad) #在这里给variable注册了 hook
out.backward()
print('*'*20)
print("x.grad", x.grad)
print(mod.weight.grad)

实验二输出
********************
x.grad tensor([[ 0.0091, -0.0202, -0.0334]]) # 增加了0.1倍
tensor([[0.1000, 0.2000, 0.3000]]) #  增加了0.1倍

实验二对照实验：无hook
x=Variable(torch.FloatTensor([[1, 2, 3]]),requires_grad=True)
mod=nn.Linear(3, 1, bias=False)
# mod.register_backward_hook(bh) # 在这里给module注册了backward hook

out=mod(x)
# out.register_hook(lambda grad: 0.1*grad) #在这里给variable注册了 hook
out.backward()
print('*'*20)
print("x.grad", x.grad)
print(mod.weight.grad)

实验二对照实验输出
********************
x.grad tensor([[0.5737, 0.3620, 0.4793]]) # 正常梯度
tensor([[1., 2., 3.]])

3 module级别的hook

实验三：module级别的register_hook函数，影响了整个模型的grad
import torch
from torch.autograd import Variable
from torch.nn import Parameter
import torch.nn as nn
import math
def bh(m,gi,go):
    print("Grad Input")
    print(gi)
    print("Grad Output")
    print(go)
    return gi[0]*0,gi[1]*0
x=Variable(torch.FloatTensor([[1, 2, 3]]),requires_grad=True)
mod=nn.Linear(3, 1, bias=False)
mod.register_backward_hook(bh) # 在这里给module注册了backward hook

out=mod(x)
#out.register_hook(lambda grad: 0.1*grad) #在这里给variable注册了 hook
out.backward()
print('*'*20)
print("x.grad", x.grad)
print(mod.weight.grad)

实验三输出：
Grad Input
(tensor([[ 0.4894, -0.0671,  0.1403]]), tensor([[1.],
        [2.],
        [3.]]))
Grad Output
(tensor([[1.]]),)
********************
x.grad tensor([[0., -0., 0.]]) # 梯度为0 ，因为bh函数的设置：return gi[0]*0,gi[1]*0
tensor([[0., 0., 0.]])

以上实验影响了梯度的反向传播，也可以不影响：

实验三对照实验
def bh(m,gi,go):
    print("Grad Input")
    print(gi)
    print("Grad Output")
    print(go)

x=Variable(torch.FloatTensor([[1, 2, 3]]),requires_grad=True)
mod=nn.Linear(3, 1, bias=False)
mod.register_backward_hook(bh) # 在这里给module注册了backward hook

out=mod(x)
out.backward()
print('*'*20)
print("x.grad", x.grad)
print(mod.weight.grad)

Grad Input
(tensor([[-0.0587,  0.2710, -0.1430]]), tensor([[1.],
        [2.],
        [3.]]))
Grad Output
(tensor([[1.]]),)
********************
x.grad tensor([[-0.0587,  0.2710, -0.1430]])
tensor([[1., 2., 3.]])

此外，似乎module级别的hook在使用的时候还要注意这个module不能是一个：
在module上注册一个bachward hook。此方法目前只能用在Module上，不能用在Container上，当Module的forward函数中只有一个Function的时候，称为Module，如果Module包含其它Module，称之为Container，见
https://blog.csdn.net/u012436149/article/details/69230136?utm_medium=distribute.pc_relevant_t0.none-task-blog-BlogCommendFromMachineLearnPai2-1.channel_param&depth_1-utm_source=distribute.pc_relevant_t0.none-task-blog-BlogCommendFromMachineLearnPai2-1.channel_param