PyTorch中Forward与Backward Hook的代码级工作原理及调用机制问询

一、钩子触发的参数传递机制 #

你提到的module.register_forward_hook(self.get_attention)这种调用，本质是Python回调函数的典型应用：

• 注册时只传函数对象，是因为PyTorch框架内部已经约定好了钩子触发时要传递的参数格式，会在合适的时机自动把参数喂给你的钩子函数。
• 比如前向钩子的函数必须遵循固定签名：hook_fn(module, input, output)，其中module就是你注册钩子的那个层（比如上面的self.linear），input是该层的输入张量，output是该层的输出张量。反向钩子的签名则是hook_fn(module, grad_input, grad_output)，对应梯度相关的参数。
• 简单说，你把钩子函数“挂”到模块的回调列表里，当模块完成前向/反向传播时，PyTorch会主动调用这个函数，并把预设好的参数传进去，不需要你手动处理参数传递。

举个实际代码例子：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 2)
        # 给linear层注册前向钩子
        self.linear.register_forward_hook(self.get_attention)
    
    def get_attention(self, module, input, output):
        # 这里的参数是PyTorch自动传进来的
        print(f"当前层：{module.__class__.__name__}")
        print(f"输入形状：{input[0].shape}")
        print(f"输出形状：{output.shape}")
    
    def forward(self, x):
        return self.linear(x)

model = MyModel()
input_tensor = torch.randn(3, 10)
output = model(input_tensor)

运行这段代码，前向传播时get_attention会自动触发，参数都是PyTorch帮你传好的。

二、Forward/Backward Hook的核心运作逻辑 #

前向钩子（Forward Hook） #

• 注册时机：模型初始化后、推理/训练前完成注册即可
• 触发时机：对应模块执行完前向传播计算，即将返回输出的时候
• 用途：提取中间层特征、调试模型输出、甚至修改中间输出（比如做特征增强）

反向钩子（Backward Hook） #

• 注册时机：反向传播开始前完成注册（一般和前向钩子一起注册）
• 触发时机：对应模块完成反向梯度计算，梯度即将向上一层回传的时候
• 用途：分析层的梯度分布、修改梯度（比如梯度裁剪）、排查梯度消失/爆炸问题

三、关于反向钩子历史Bug的说明 #

早期PyTorch（1.0版本之前）确实存在反向钩子的一些问题，比如部分模块的grad_input返回值不符合预期、钩子偶尔不触发等。但从PyTorch 1.0之后的稳定版本开始，这些问题已经陆续修复。目前主流的PyTorch版本（1.10及以上）中，反向钩子的逻辑已经非常稳定，只要你严格按照规定的函数签名编写钩子，基本不会遇到历史遗留的Bug。如果使用中遇到异常，优先检查钩子函数的签名是否正确，或者是否在反向传播前正确注册了钩子。

内容的提问来源于stack exchange，提问作者Jose Ramon