先说结论：1. hook机制主要有三种表现形式，即register_hook()、register_forward_hook()、register_backward_hook()。register_hook()主要应用在tensor层面，后两者主要用在module层面。
2. 以module层面为例，通俗来讲，hook的作用就是就是在不改变module定义源码的前提下，获取模型某一层的属性。register_forward_hook()主要是获取前向传播的一些属性，例如某一层的name、某一层的input feature map或output feature map等；register_backward_hook()主要是获取反向传播的属性，主要是梯度。
一、Hook机制提出的意义在Pytorch中，模型在计算过程中会自动舍弃计算的中间结果，所以想要获取这些数值就需要使用hook函数。例如，对于复合函数f[g(x)]而言，若要求f关于的导数，根据链式法则必然要先求f关于g的导数，但是在PyTorch中不会保留f关于g的导数这一中间计算结果。因此，在不改变网络的定义代码的前提下，也不需要在forward函数中return某个感兴趣层的输出，利用hook来获取某一层的属性。
二、Hook机制的应用举例<1>使用注意事项
register_forward_hook()函数必须在forward()函数调用之后被使用，因为这个函数源码注释显示这个函数“ it will not have effect on forward since this is called after :func:`forward` is called”。
<2>实验代码①
使用方式为module.register_forward_hook(hook_fn)，其中hook_fn的声明为:
hook_fn(module, input, output) -> Tensor or None
函数的输入分别为module模块，输入，输出。hook可以修改输出值，使用方法类似于Tensor的Hook。
def hook_forward_fn(module, input, output):
print(f"module: {module}")
print(f"input: {input}")
print(f"output: {output}")
class CBR(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3, 1, 1)
self.bn = nn.BatchNorm2d(1)
self.conv.register_forward_hook(hook_forward_fn)
self.bn.register_forward_hook(hook_forward_fn)
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
return torch.relu(x)
net = CBR()
x = torch.rand(1, 3, 2, 2, requires_grad=True)
y = net(x).mean()
y.backward()
输出结果为：
module: Conv2d(3, 1, kernel_size=(1, 1), stride=(1, 1))
input: (tensor([[[[0.1824, 0.0140],
[0.3478, 0.8911]],
[[0.3780, 0.3915],
[0.3625, 0.1337]],
[[0.5159, 0.2859],
[0.5742, 0.0458]]]], requires_grad=True),)
output: tensor([[[[-0.3373, -0.2838],
[-0.3712, -0.3378]]]], grad_fn=)
module: BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
input: (tensor([[[[-0.3373, -0.2838],
[-0.3712, -0.3378]]]], grad_fn=),)
output: tensor([[[[-0.1507, 1.5480],
[-1.2299, -0.1675]]]], grad_fn=)
<3>实验代码②
使用方式为module.register_backward_hook)hook_fn)，其中hook_fn的声明为:
hook_fn(module, grad_input, grad_output) -> Tensor or None
函数的输入分别为module模块，输入端的grad，输出端的grad，这里称呼的输入与输出是站在前向传播的角度的。如果模块有多个输入与输出的话， 其grad_input和grad_output可以是tuple类型。 与forward不同的是，backward传播的时候，不仅反向传递input和output的grad，还会传递模块Parameter的grad，因此比如fc模块，其grad_input 是一个三元组的tuple，（对bias的梯度，对输入的梯度，对w的梯度）；conv模块grad_input也是一个三元组tuple，为（对输入的梯度，对w的梯度，对bias的梯度）。如果想改变梯度，可以对grad_output进行操作。
def hook_backward_fn(module, grad_input, grad_output):
print(f"module: {module}")
print(f"grad_output: {grad_output}")
print(f"grad_input: {grad_input}")
print("*"*20)
class fc(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(1, 2)
self.fc2 = nn.Linear(2, 3)
self.fc1.register_backward_hook(hook_backward_fn)
self.fc2.register_backward_hook(hook_backward_fn)
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
return x
x = torch.rand(2, 1, requires_grad=True)
net = fc()
y = net(x).mean()
y.backward()
输出结果为：
module: Linear(in_features=2, out_features=3, bias=True)
grad_output: (tensor([[0.1667, 0.1667, 0.1667],
[0.1667, 0.1667, 0.1667]]),)
grad_input: (tensor([0.3333, 0.3333, 0.3333]), tensor([[-0.0926, 0.0243],
[-0.0926, 0.0243]]), tensor([[-0.1773, -0.1773, -0.1773],
[-0.0068, -0.0068, -0.0068]]))
********************
module: Linear(in_features=1, out_features=2, bias=True)
grad_output: (tensor([[-0.0926, 0.0243],
[-0.0926, 0.0243]]),)
grad_input: (tensor([-0.1853, 0.0486]), tensor([[0.0073],
[0.0073]]), tensor([[-0.1057, 0.0278]]))
********************
<4>实验代码③
这里展示的是比较常见的hook_fn的定义方式：
import torch
import torch.nn as nn
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16*5*5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
out = F.relu(self.conv1(x)) #1
out = F.max_pool2d(out, 2) #2
out = F.relu(self.conv2(out)) #3
out = F.max_pool2d(out, 2)
out = out.view(out.size(0), -1)
out = F.relu(self.fc1(out))
out = F.relu(self.fc2(out))
out = self.fc3(out)
return out
features = []
def hook(module, input, output):
# module: model.conv2
# input :in forward function [#2]
# output:is [#3 self.conv2(out)]
features.append(output.clone().detach())
# output is saved in a list
net = LeNet() ## 模型实例化
x = torch.randn(2, 3, 32, 32) ## input
handle = net.conv2.register_forward_hook(hook) ## 获取整个Lenet模型 conv2的中间结果
y = net(x) ## 获取的是 关于 input x 的 conv2 结果
print(features[0].size()) # 即 [#3 self.conv2(out)]
handle.remove() ## hook删除
<5>实验代码④
若要绘制热力图，首先要获取梯度。下面附两个绘制热力图的参考链接：
PaddlePaddle热力图绘制
Matplotlib调用imshow()函数绘制热图
三、Reference实验代码参考来源：
代码参考1
代码参考2
