Pytorch中的钩子函数Hook函数

1. 为什么要使用Hook函数？

因为中间变量完成了反向传播后就自动释放了，因此无法读出存储的梯度。

2. 有什么样的Hook函数

• torch.autograd.Variable.register_hook

import torch

def hook_fn(grad):
    print("Gradient:", grad)

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x + 2
z = torch.mean(torch.pow(y, 2))

y.register_hook(hook_fn)

z.backward()
1
2
3
4
5
6
7
8
9
10
11
12

在这个例子中，我们在变量 y 上注册了钩子函数 hook_fn。当调用 z.backward() 进行反向传播计算梯度时，钩子函数 hook_fn 会被自动调用，并打印出相应的梯度值。

• torch.nn.Module.register_backward_hook

import torch

def hook_fn(module, grad_input, grad_output):
    # 提取中间层的梯度
    intermediate_gradient = grad_input[0]
    # 对梯度进行处理或记录操作
    # ...

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 16, kernel_size=3)
        self.conv2 = torch.nn.Conv2d(16, 32, kernel_size=3)
        # ...

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        # ...
        return x

model = MyModel()

# 在中间层conv1上注册钩子函数
model.conv1.register_backward_hook(hook_fn)

# 在中间层conv2上注册钩子函数
model.conv2.register_backward_hook(hook_fn)

input_data = torch.randn(1, 3, 32, 32)
output = model(input_data)

loss = output.sum()
loss.backward()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

在上述示例中，我们在模型的 conv1 和 conv2 层上分别注册了钩子函数 hook_fn。当模型进行反向传播时，钩子函数将分别捕获这两个中间层的梯度。

• torch.nn.Module.register_forward_hook

import torch
import matplotlib.pyplot as plt

def visualize_feature_map(module, input, output):
    # 可视化输出特征图
    feature_map = output.detach().squeeze()
    plt.imshow(feature_map, cmap='gray')
    plt.show()

# 创建一个模块
conv = torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)

# 注册钩子函数，在前向传播时可视化输出特征图
conv.register_forward_hook(visualize_feature_map)

# 输入数据并进行前向传播
input_data = torch.randn(1, 3, 32, 32)
output = conv(input_data)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

在这个例子中，我们创建了一个卷积模块conv，然后注册了一个钩子函数visualize_feature_map。该钩子函数在模块的前向传播过程中被调用，并可视化输出的特征图。

参考：https://www.zhihu.com/question/61044004/answer/183682138