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()

在这个例子中，我们在变量 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()

在上述示例中，我们在模型的 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)

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

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