1. Hook在自动求导链路中的位置先搞懂它到底在“钩”什么1.1 一次实际排查引发的思考有一次我在调一个图像分割模型训练到第40轮左右loss突然出现NaN。常规操作是先检查输入有没有异常、学习率是否过大、loss有没有除零。这些都没问题。后来我把中间层梯度打印出来发现在某个BatchNorm层之后梯度的分布已经彻底崩了——不是所有参数都变成NaN而是一条通道的梯度比其他通道大了几个数量级。那时候我第一反应是修改forward函数把中间层输出手动保存下来再在backward之后去访问.grad。但这样一来,不仅要改前向代码还要小心loss.backward()之后的梯度是否已经被释放操作起来很别扭。后来真正解决问题的是Hook。在PyTorch里Hook就像在自动求导和图执行链条上挂一个“探针”或“拦截器”前向执行到某个模块时触发一次反向传播经过某个张量或模块时再触发一次。你可以在不修改模型原有forward逻辑的情况下读取、打印、修改甚至替换梯度。它不是为了替代手动debug而存在的而是为了让你在模型结构不变的前提下获得对中间状态的控制权。1.2 “钩子”到底挂在图的哪个位置想用好Hook必须先理解PyTorch的自动求导图。每次前向运算PyTorch都会动态构建一个由Tensor和Function节点组成的计算图。张量是数据节点Function是运算节点。调用loss.backward()时梯度从输出的Tensor开始沿着计算图反向传播每经过一个Function就会算出它对输入和参数的局部梯度再继续往前传。Hook的触发点就在这个反向传播路径上Tensor.register_hook()挂在这个Tensor上当这个Tensor的梯度被计算出来时触发。注意是“梯度已经算好、准备继续向前传”的时刻。Module.register_forward_hook()模块的forward()执行完之后触发输入输出都拿得到。Module.register_full_backward_hook()模块的反向传播相关梯度算完之后触发拿的是模块输入和输出的梯度。Module.register_forward_pre_hook()模块的forward()执行之前触发可以修改输入。把这条链路在脑中过一遍前向执行到某个Module时forward hook先于或后于模块的输出产生反向传播经过某个Tensor时tensor hook先于梯度继续向前传递而触发。这些钩子不会改变计算图的拓扑只是插在某个节点上做“旁路操作”。理解这一点后面所有的用法都不会偏。1.3 什么时候该用Hook什么时候不该用不是所有中间变量都值得用Hook。如果你的需求只是“看一眼”某个层的输出那在forward里顺手print(x.shape)也可以但这样会污染前向代码尤其是复现他人模型时你不想在源码里到处加print。Hook的价值在于分离关注点模型结构保持干净调试和分析逻辑挂在外面。以下场景特别适合用Hook一次训练循环里同时观察多个层的梯度统计不想逐个写.grad判断。需要修改梯度比如梯度裁剪只针对某一层、对某类样本的梯度做缩放。想提取中间层特征做可视化、风格迁移或对比学习。需要研究梯度消失/爆炸定位NaN出现在计算图的哪一段。不太适合用Hook的情况也很明确如果只是修改模型的最终输出或计算自定义loss那直接在forward或loss函数里做就好了用Hook反而增加理解成本。Hook更适合“旁路观察”和“局部干预”。这就像在电路板上测信号不是每个点都需要烙铁焊接有些地方用示波器探头夹一下就行——Hook就是那个探头。2. Tensor级与Module级Hook两套接口的分工与误用排查2.1register_hookTensor特有的反向拦截先看最常用的Tensor Hook。它的使用非常简单import torch t torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) def print_grad(grad): print(grad:, grad) return grad * 2 # 返回的梯度会继续向后传播 h t.register_hook(print_grad) loss (t * t).sum() loss.backward() # 输出 grad: tensor([2., 4., 6.]) print(t.grad) # tensor([4., 8., 12.])这个例子包含了两个关键点register_hook接收的函数必须接收一个梯度张量并返回一个梯度张量。返回值为None时不修改梯度只做观察。如果返回了新的梯度后续传播使用的就是被替换后的梯度。上面t * t对t的梯度是2*t也就是[2,4,6]hook把它翻倍变成[4,8,12]最后t.grad保存的就是修改后的值。我在实际项目里经常利用这一点做“按样本调整梯度”。比如在训练某个模拟项目X时为了让模型对困难样本更敏感我会在最后一个embedding张量上挂hook根据该样本的feature norm对梯度做缩放。这种逻辑写进训练循环里会很繁琐但用hook只挂一次就行。2.2 Module级Hook拿到输入、输出和梯度Module级Hook比Tensor Hook更容易理解因为它正对着你的模型层。下面这段代码可以同时观察某个卷积层的输入输出形状以及反向传播时输入梯度的分布import torch import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, 3) self.relu nn.ReLU() self.fc nn.Linear(16 * 6 * 6, 10) def forward(self, x): x self.conv1(x) x self.relu(x) x x.flatten(1) x self.fc(x) return x net SimpleNet() def forward_hook(module, input, output): print(f[forward] {module.__class__.__name__} input: {input[0].shape}, output: {output.shape}) def backward_hook(module, grad_input, grad_output): print(f[backward] {module.__class__.__name__}) for i, g in enumerate(grad_input): if g is not None: print(f grad_input[{i}]: {g.shape}) for i, g in enumerate(grad_output): if g is not None: print(f grad_output[{i}]: {g.shape}) h1 net.conv1.register_forward_hook(forward_hook) h2 net.conv1.register_full_backward_hook(backward_hook) x torch.randn(2, 3, 8, 8) out net(x) loss out.sum() loss.backward()注意我使用的是register_full_backward_hook而不是旧版的register_backward_hook。旧版接口在很多新版本中已经被标记为不推荐因为它的grad_input可能不是完整的输入梯度而且默认情况下模块输出要求非多个Tensor否则可能不触发。新版full_backward_hook行为更直观在模块的反向传播计算完成后触发传入grad_input模块输入的梯度和grad_output模块输出的梯度。注意这里的命名容易混淆grad_input指的是“这个模块的输入所接收到的梯度”不是模型输入grad_output指的是“这个模块的输出所接收到的梯度”。2.3 实战中容易搞混的四种状态很多人一开始会把Module hook拿到的input/output和grad_input/grad_output搞混然后debug半天。我整理了一个对应关系表钩子类型触发时机第一参数第二参数主要用途forward_pre_hook前向执行前模块输入无修改输入、注入扰动forward_hook前向执行后模块输入模块输出提取特征图、记录形状full_backward_hook反向传播后模块输入梯度元组模块输出梯度元组梯度诊断、梯度缩放tensor.register_hookTensor梯度算出后待传梯度无单点梯度修改、梯度裁剪其中full_backward_hook的grad_input元组里有可能会出现None。比如ReLU层它对负区间的导数为0对应位置的梯度就是None。你写统计时要先过滤None否则直接torch.max(grad_input[0])会报错。这是我第一次写Hook时踩到的第一个坑后面会细说。2.4 如何判断钩子有没有生效一个很实用的技巧是临时用一个“打印型hook”验证代码逻辑。在怀疑某个钩子没有触发时不要直接看结果先打印个标记def tmp_hook(module, input, output): print(hook fired:, module.__class__.__name__)如果什么都没打印大概率是模块根本没有被执行而不是钩子没挂上。比如在一个nn.Sequential里定义了某个层但forward没有调用它钩子自然不触发。再比如模型被torch.compile或JIT化后某些钩子行为会改变这在PyTorch 2.x版本下尤其要注意。还有一个经常被忽视的点钩子对象需要被保存。如果你的register_*_hook()返回值没有被任何变量引用Python的垃圾回收可能把它回收掉钩子随着失效。在交互式环境Jupyter里更容易遇到这种情况。所以一定要写handle net.conv1.register_forward_hook(...)并且在不需要的时候调用handle.remove()。3. 进阶实战一梯度修改、裁剪与回填的标准套路3.1 场景只想对某一层做梯度裁剪torch.nn.utils.clip_grad_norm_是全局梯度裁剪但有时你只想裁剪某个子模块的梯度。比如GAN训练中生成器的某个残差块梯度异常而判别器正常又比如混合精度训练中某些层梯度爆炸其它层没问题。全局裁剪会拖累所有层的学习速度局部裁剪更精准。用Tensor hook实现局部裁剪非常干净import torch import torch.nn as nn from torch.nn.utils import clip_grad_norm_ def make_clip_hook(max_norm): def hook(grad): return clip_grad_norm_([grad], max_norm)[0] # 返回裁剪后的梯度 return hook model SomeModel() handle model.specific_layer.weight.register_hook(make_clip_hook(1.0))注意clip_grad_norm_接收一个参数列表并返回裁剪后的总范数。我在这里用了[grad]把当前梯度包成列表原地裁剪后取回第一个元素作为新梯度返回。这是一个“少有人写但确实可行”的技巧。更稳妥的写法是手动计算def clip_hook(grad): max_norm 1.0 total_norm grad.norm().item() if total_norm max_norm: return grad * (max_norm / (total_norm 1e-6)) return grad这个手动版本的好处是你能在打印日志时看到总范数方便确认hook确实在工作。3.2 场景梯度衰减与正负号干预有些损失函数设计想让某个分支的梯度以一定比例衰减比如让多任务学习中的某个任务“慢一点学”。最笨的方法是在loss里加权重但这样会同时影响该任务所有相关层的梯度。如果只想控制从某个特定张量往后传的梯度大小那Tensor hook是最精确的位置。举一个例子模拟项目Y中我让一个共享encoder的特征向量经过两个head分别做分类和回归。回归任务常常因为样本量少而学得太快影响共享特征。我就在回归head的输入特征上挂了一个缩放hookdef scale_grad_by(grad, factor0.1): return grad * factor feat shared_encoder(x) handle feat.register_hook(lambda grad: grad * 0.1) logits_reg reg_head(feat)这样分类任务回传的梯度不受影响回归任务对共享encoder的贡献被主动削弱。其实这与在loss中乘系数不同后者会把回归loss的所有反向传播梯度同步放大或缩小但不会改变回归head内部参数的相对梯度方向而hook挂在共享输入上只影响“经过这个张量的梯度”有更精确的控制边界。3.3 场景处理retain_graph都解决不了的梯度释放问题默认情况下PyTorch在backward()后会释放计算图对于非leaf Tensor的梯度会被丢弃方便下一次迭代。如果你在训练循环里做了两次backward比如GAN要先更新判别器再更新生成器很多时候需要设置retain_graphTrue这时内存压力会比较大。利用hook可以在某些情况下避免重复backward。举个例子如果两个loss共享同一个骨干网络但你想让第一个loss只更新head A第二个loss更新head B和骨干理论上可以用“stop-gradient”类的操作但更灵活的是在第一个loss反向传播前对head A输出后的共享特征张量挂一个将梯度置零的hookfeat backbone(x) feat.register_hook(lambda grad: torch.zeros_like(grad)) loss1 head_A(feat).sum() loss1.backward() # head_A更新backbone的梯度被清零不影响后续其他loss反向传播这个方案有争议因为零梯度仍然会继续向前传播徒增计算量。如果想真正截断应该用feat.detach()或torch.where(True, feat, 0)这类方式。但零梯度hook在你要“保留计算图结构不重新forward”的场景下有独特价值通过hook先清零再backward之后的第二次backward可以复用计算图但又不会重复累加骨干梯度。当然实际使用时需要对自己的梯度累加逻辑非常清楚我建议先用简单模型验证再上大模型。3.4 场景把梯度“回填”到参数上有时你不想在backward()之后遍历named_parameters()处理梯度而是希望梯度在产生的同时就被“改造”成最终形态。比如你想统计梯度的一阶矩和二阶矩或者把梯度量化后再更新这时可以用hook把梯度保存到一个全局字典里供优化器自定义逻辑使用。grad_storage {} def store_grad(name): def hook(grad): grad_storage[name] grad.detach().clone() return grad return hook handles [] for name, param in model.named_parameters(): if param.requires_grad: handles.append(param.register_hook(store_grad(name)))注意这里必须是detach().clone()。如果只写clone()保存的张量仍与计算图有连接多次backward后可能造成内存累积。另一方面如果你想在hook里检查梯度并在异常时“回填”梯度直接修改返回的grad即可。很多人误以为hook只能“看”不能“改”其实返回值就是用于继续传播的梯度拥有完全控制权。4. 进阶实战二特征图与梯度的诊断统计定位训练异常4.1 用hook采集多层的梯度范数定位NaN源头训练遇到NaN时最痛苦的是不知道NaN是从哪一层先冒出来的。一个高效的排查方案是在每一个卷积层注册full_backward_hook计算每个模块输出梯度的范数和均值实时打印。下面是我常用的诊断代码骨架import torch import torch.nn as nn def make_grad_monitor(module_name, log_freq50): def backward_hook(module, grad_input, grad_output): if grad_output is None: return grad_t grad_output[0] if grad_t is None: return fn torch.isnan(grad_t).float().mean().item() norm grad_t.norm().item() if fn 0 or norm 1e10: print(f[nan/explode] {module_name}: nan_ratio{fn:.6f}, norm{norm:.6g}) return backward_hook for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear, nn.BatchNorm2d)): module.register_full_backward_hook(make_grad_monitor(name))大多数情况下NaN并不是所有层同一时刻一起爆掉的而是从某一层开始梯度变成NaN然后顺着反向传播传播到前面的层。通过这种逐层监控你能看到第一个报警的模块名那通常就是问题根源。我在模拟项目X里曾经用这种方式定位到是transformer里某个attention权重计算溢出了——QK之后没有缩放数值直接跑到几千softmax之后的梯度出现了Inf。4.2 用forward hook提取特征图做中间激活的可视化很多同学以为提取特征图必须在网络内部保存或者要重新forward一遍。其实用forward hook完全可以在模型forward完成之后从外部拿到任意层的输出。activation {} def save_activation(name): def hook(module, input, output): activation[name] output.detach() return hook model.conv1.register_forward_hook(save_activation(conv1)) model.layer2.register_forward_hook(save_activation(layer2)) _ model(x) feat activation[conv1] # (B, C, H, W)这里有几个工程化小技巧一定要detach()否则activation字典会持有整张计算图内存可能爆炸。如果只想存最后一次forward的结果记得在每次forward前activation.clear()。如果在一批样本上提取特征做聚类最好不要把特征直接存到增量的list里而是提前分配一个大Tensor用索引写入避免反复append带来的显存碎片。拿到特征图后最常见的处理是对Channel维做归一化再保存成图片import torch.nn.functional as F feat_min feat.min() feat_max feat.max() feat_norm (feat - feat_min) / (feat_max - feat_min 1e-8) grid torchvision.utils.make_grid(feat_norm[0].unsqueeze(0), nrow8) torchvision.utils.save_image(grid, feat.png)4.3 用hook统计梯度分布而不是只看范数范数只能告诉你“大还是小”不能告诉你“梯度里有没有异常模式”。我更推荐在关键层统计梯度直方图的分位数比如打印p50/p95/p99。因为有时候范数正常但99%的维度已经消失只有1%的维度在传梯度这种情况对训练是灾难。用hook实现分位数统计很简单def quantile_hook(module_name): def hook(grad): flat grad.flatten().abs() if flat.numel() 0: q torch.quantile(flat, torch.tensor([0.5, 0.95, 0.99]).to(grad.device)) print(f{module_name}: p50{q[0]:.4g}, p95{q[1]:.4g}, p99{q[2]:.4g}) return grad return hook把这个hook挂到容易出现梯度消失的层训练前期观察数值变化能第一时间发现初始化是否合适。比如一个Sigmoid激活块的p99长期小于1e-8基本可以断定这一块进入饱和区考虑换激活函数或加残差。4.4 自动化把hook采集的数据汇总成训练日志Hook的代码容易散落在训练脚本各处我建议封装成一个“诊断器”类集中管理handle和采集结果class GradientInspector: def __init__(self, model, module_types(nn.Conv2d, nn.Linear)): self.handles [] self.records {} for name, module in model.named_modules(): if isinstance(module, module_types): self.handles.append( module.register_full_backward_hook(self._make_hook(name)) ) def _make_hook(self, name): def hook(module, grad_input, grad_output): g grad_output[0] if g is not None: self.records[name] { norm: g.norm().item(), mean: g.mean().item(), absmax: g.abs().max().item(), } return hook def get_report(self): return self.records def clear(self): self.records.clear() def remove(self): for h in self.handles: h.remove()在训练循环的某个频率下调用get_report()记录到TensorBoard或wandb基本就是一条轻量级的梯度监控流水线。不需要再侵入模型内部的forward代码。这也是“进阶篇”该有的样子不只写得出hook还要把它变成可复用的工具。5. 我最想提醒的Hook坑inplace、作用域顺序与remove机制5.1 inplace操作会让hook的梯度指向被覆盖的张量我在做模型压缩实验时为了省显存在某个层的输出后面直接用了relu_(x)这类inplace操作。结果发现hook拿到的输出和张量本身总是对不上某些统计结果看起来“差一点点”。问题根源是inplace操作会直接在原内存上修改张量而计算图的Function节点记录的仍然是旧Tensor对象引用。反向传播时hook从某个节点拿到的梯度可能对应的是修改前还是修改后的版本行为取决于具体实现完全不可预测。尽量在hook涉及的路径上避免inplace操作。比如ReLU用nn.ReLU()而不是torch.relu_(x)在Transformer的残差连接处不要用x attn(x)改为x x attn(x)。现代PyTorch对inplace操作支持得越来越好但在hook这种“旁路读取”场景中inplace会带来诡异的别名问题排查成本非常高。如果为了显存必须用inplace那就要保证hook只用来观察不要用来修改梯度。5.2 多个hook的顺序不是你想的那样在同一个Tensor或Module上挂了多个hook触发顺序如何PyTorch的文档并没有给非常严格的全局顺序承诺实际经验是多个tensor.register_hook会按注册顺序执行前一个返回的梯度作为后一个的输入Module多个forward hook也会按注册顺序依次执行。但你不能因此就写“依赖顺序”的业务逻辑官方没有把顺序当作公开契约。跨Tensor和跨Module的hook顺序就更难预测了因为反向传播本身是图的拓扑序不同分支的先后并不固定。这个坑在实际中非常常见。比如你想先缩放梯度再做梯度裁剪如果你把它们挂成两个hook依赖执行顺序是不稳定的。稳妥做法是只写一个hook内部自己做两步操作def combined_hook(grad): grad grad * 0.1 grad clip_by_norm(grad) return grad5.3 hook在nn.Sequential和torch.compile下的行为差异nn.Sequential本身是一个Module但hook你通常想挂在里面的子模块上这没问题。真正麻烦的是PyTorch 2.x推出torch.compile后模型会被重写或融合很多子模块在运行时并不以Python对象形式存在此时register_forward_hook可能不触发因为前向已经被编译成融合的C算子。你在调试一个compile过的模型时如果某个hook静默失效不要怀疑是自己写错了先试试在非compile模式下运行。我的经验是训练和推理用compile加速调试和Hook分析关掉compile。两者分开既保速度又保可观测性。如果你必须要在compile模式下挂hook可以试试用torch._dynamo.config.export_mode或者对指定模块加torch.compiler.disable装饰让那个模块用原始Python执行但这是比较偏门的做法新手阶段不建议深挖。5.4 handle.remove()与生命周期管理我在2.4节提过handle要保存这里再展开说。很多人写论文代码时在训练循环里注册hook循环结束后不remove然后反复实例化模型做多组实验结果发现hook越来越多内存越占越高梯度统计也越来越乱。原因就是每次注册产生的handle被丢弃但注册过的hook不会被垃圾回收自动清理除非对应的模型对象被销毁。建议写一个统一的hook管理类所有hook的handle都放进去训练循环结束时统一调用remove()。比如hook_handles [] def add_hook(module, hook_fn): h module.register_forward_hook(hook_fn) hook_handles.append(h) def clear_hooks(): for h in hook_handles: h.remove() hook_handles.clear()这个习惯能帮你避免很多莫名其妙的“它为什么影响了另一个实验”的问题。另外不要在临时函数里注册hook却不返回handle比如def setup(): model.conv1.register_forward_hook(fn) # 不保存handle危险这种写法在交互式环境里经常让hook残留影响后续所有输入的前向输出。5.5detach、requires_grad与hook返回值的关系最后一个容易踩坑的点是即使你在hook里返回了新的梯度原Tensor的.grad也只会在它本身是leaf Tensor时被保存。对非leaf Tensor来说hook处理后返回的梯度会继续向前传播但不会永久保存在该Tensor的.grad属性里。很多人在中间层张量上挂hook期望返回梯度后这个张量有.grad可以读结果发现是None误以为hook没生效。其实hook生效了只是非leaf Tensor默认不保留梯度。如果你希望在hook里既修改梯度又保存梯度必须像前面那样grad.detach().clone()自己存到外部数据结构里。不要试图通过设置tensor.retain_grad()来同时兼得——和hook一起使用时要先明确retain_grad()的作用是让非leaf Tensor保存反向梯度它和hook的返回修改不冲突但会让显存占用上升。6. 收尾一个从调试到落地的Hook应用模板最后分享一个我自己的做法。我通常会把Hook代码独立成一个模块不写在训练脚本里方便后续复用到其它项目。整个“调试型Hook”大概分三层第一层是“记录层”只收集不修改。比如forward hook提取特征图、backward hook统计梯度分位数。 第二层是“干预层”根据记录层的统计结果动态决定是否修改梯度。比如检测到梯度范数超过阈值就放缩梯度。 第三层是“清理层”训练结束后统一remove所有handle输出统计报告。把这个模板写到你的项目里会少走很多弯路class DebugHookManager: def __init__(self, model): self.model model self.handles [] self.storage {} def register_forward(self, name, module): def hook(module, input, output): self.storage[fout:{name}] output.detach().clone() self.handles.append(module.register_forward_hook(hook)) def register_backward(self, name, module): def hook(module, grad_input, grad_output): g grad_output[0] if g is not None: self.storage[fgrad:{name}] g.detach().clone() self.handles.append(module.register_full_backward_hook(hook)) def remove(self): for h in self.handles: h.remove() self.handles.clear()当然这个模板只是一个起点。Hook真正的威力在于它允许你把“观察逻辑”和“模型结构”解耦。以前调试一个几十层的模型我要在每层后面写一堆print和assert改代码时心惊胆战生怕漏删。现在写好在外部模型完全不用动。训练异常时挂上去定位完移除干净利落。如果你刚接触Hook我的建议是从一个最简单的实验开始在一个两层线性模型上挂tensor.register_hook手动算一下梯度验证返回值和原始梯度的关系。之后再过渡到真实模型。不要一上来就在几十层的网络上搞花活因为Hook的调试本身也有成本局部出错很难定位。先把Tensor、Module两套接口的触发时机吃透后面的进阶用法就都是顺水推舟了。 
长河君 · 2026/10/18 21:21:01 · 阅读 1.2k PyTorch Hook机制详解：从梯度调试到特征提取的实战指南
1. Hook在自动求导链路中的位置先搞懂它到底在“钩”什么1.1 一次实际排查引发的思考有一次我在调一个图像分割模型训练到第40轮左右loss突然出现NaN。常规操作是先检查输入有没有异常、学习率是否过大、loss有没有除零。这些都没问题。后来我把中间层梯度打印出来发现在某个BatchNorm层之后梯度的分布已经彻底崩了——不是所有参数都变成NaN而是一条通道的梯度比其他通道大了几个数量级。那时候我第一反应是修改forward函数把中间层输出手动保存下来再在backward之后去访问.grad。但这样一来,不仅要改前向代码还要小心loss.backward()之后的梯度是否已经被释放操作起来很别扭。后来真正解决问题的是Hook。在PyTorch里Hook就像在自动求导和图执行链条上挂一个“探针”或“拦截器”前向执行到某个模块时触发一次反向传播经过某个张量或模块时再触发一次。你可以在不修改模型原有forward逻辑的情况下读取、打印、修改甚至替换梯度。它不是为了替代手动debug而存在的而是为了让你在模型结构不变的前提下获得对中间状态的控制权。1.2 “钩子”到底挂在图的哪个位置想用好Hook必须先理解PyTorch的自动求导图。每次前向运算PyTorch都会动态构建一个由Tensor和Function节点组成的计算图。张量是数据节点Function是运算节点。调用loss.backward()时梯度从输出的Tensor开始沿着计算图反向传播每经过一个Function就会算出它对输入和参数的局部梯度再继续往前传。Hook的触发点就在这个反向传播路径上Tensor.register_hook()挂在这个Tensor上当这个Tensor的梯度被计算出来时触发。注意是“梯度已经算好、准备继续向前传”的时刻。Module.register_forward_hook()模块的forward()执行完之后触发输入输出都拿得到。Module.register_full_backward_hook()模块的反向传播相关梯度算完之后触发拿的是模块输入和输出的梯度。Module.register_forward_pre_hook()模块的forward()执行之前触发可以修改输入。把这条链路在脑中过一遍前向执行到某个Module时forward hook先于或后于模块的输出产生反向传播经过某个Tensor时tensor hook先于梯度继续向前传递而触发。这些钩子不会改变计算图的拓扑只是插在某个节点上做“旁路操作”。理解这一点后面所有的用法都不会偏。1.3 什么时候该用Hook什么时候不该用不是所有中间变量都值得用Hook。如果你的需求只是“看一眼”某个层的输出那在forward里顺手print(x.shape)也可以但这样会污染前向代码尤其是复现他人模型时你不想在源码里到处加print。Hook的价值在于分离关注点模型结构保持干净调试和分析逻辑挂在外面。以下场景特别适合用Hook一次训练循环里同时观察多个层的梯度统计不想逐个写.grad判断。需要修改梯度比如梯度裁剪只针对某一层、对某类样本的梯度做缩放。想提取中间层特征做可视化、风格迁移或对比学习。需要研究梯度消失/爆炸定位NaN出现在计算图的哪一段。不太适合用Hook的情况也很明确如果只是修改模型的最终输出或计算自定义loss那直接在forward或loss函数里做就好了用Hook反而增加理解成本。Hook更适合“旁路观察”和“局部干预”。这就像在电路板上测信号不是每个点都需要烙铁焊接有些地方用示波器探头夹一下就行——Hook就是那个探头。2. Tensor级与Module级Hook两套接口的分工与误用排查2.1register_hookTensor特有的反向拦截先看最常用的Tensor Hook。它的使用非常简单import torch t torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) def print_grad(grad): print(grad:, grad) return grad * 2 # 返回的梯度会继续向后传播 h t.register_hook(print_grad) loss (t * t).sum() loss.backward() # 输出 grad: tensor([2., 4., 6.]) print(t.grad) # tensor([4., 8., 12.])这个例子包含了两个关键点register_hook接收的函数必须接收一个梯度张量并返回一个梯度张量。返回值为None时不修改梯度只做观察。如果返回了新的梯度后续传播使用的就是被替换后的梯度。上面t * t对t的梯度是2*t也就是[2,4,6]hook把它翻倍变成[4,8,12]最后t.grad保存的就是修改后的值。我在实际项目里经常利用这一点做“按样本调整梯度”。比如在训练某个模拟项目X时为了让模型对困难样本更敏感我会在最后一个embedding张量上挂hook根据该样本的feature norm对梯度做缩放。这种逻辑写进训练循环里会很繁琐但用hook只挂一次就行。2.2 Module级Hook拿到输入、输出和梯度Module级Hook比Tensor Hook更容易理解因为它正对着你的模型层。下面这段代码可以同时观察某个卷积层的输入输出形状以及反向传播时输入梯度的分布import torch import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, 3) self.relu nn.ReLU() self.fc nn.Linear(16 * 6 * 6, 10) def forward(self, x): x self.conv1(x) x self.relu(x) x x.flatten(1) x self.fc(x) return x net SimpleNet() def forward_hook(module, input, output): print(f[forward] {module.__class__.__name__} input: {input[0].shape}, output: {output.shape}) def backward_hook(module, grad_input, grad_output): print(f[backward] {module.__class__.__name__}) for i, g in enumerate(grad_input): if g is not None: print(f grad_input[{i}]: {g.shape}) for i, g in enumerate(grad_output): if g is not None: print(f grad_output[{i}]: {g.shape}) h1 net.conv1.register_forward_hook(forward_hook) h2 net.conv1.register_full_backward_hook(backward_hook) x torch.randn(2, 3, 8, 8) out net(x) loss out.sum() loss.backward()注意我使用的是register_full_backward_hook而不是旧版的register_backward_hook。旧版接口在很多新版本中已经被标记为不推荐因为它的grad_input可能不是完整的输入梯度而且默认情况下模块输出要求非多个Tensor否则可能不触发。新版full_backward_hook行为更直观在模块的反向传播计算完成后触发传入grad_input模块输入的梯度和grad_output模块输出的梯度。注意这里的命名容易混淆grad_input指的是“这个模块的输入所接收到的梯度”不是模型输入grad_output指的是“这个模块的输出所接收到的梯度”。2.3 实战中容易搞混的四种状态很多人一开始会把Module hook拿到的input/output和grad_input/grad_output搞混然后debug半天。我整理了一个对应关系表钩子类型触发时机第一参数第二参数主要用途forward_pre_hook前向执行前模块输入无修改输入、注入扰动forward_hook前向执行后模块输入模块输出提取特征图、记录形状full_backward_hook反向传播后模块输入梯度元组模块输出梯度元组梯度诊断、梯度缩放tensor.register_hookTensor梯度算出后待传梯度无单点梯度修改、梯度裁剪其中full_backward_hook的grad_input元组里有可能会出现None。比如ReLU层它对负区间的导数为0对应位置的梯度就是None。你写统计时要先过滤None否则直接torch.max(grad_input[0])会报错。这是我第一次写Hook时踩到的第一个坑后面会细说。2.4 如何判断钩子有没有生效一个很实用的技巧是临时用一个“打印型hook”验证代码逻辑。在怀疑某个钩子没有触发时不要直接看结果先打印个标记def tmp_hook(module, input, output): print(hook fired:, module.__class__.__name__)如果什么都没打印大概率是模块根本没有被执行而不是钩子没挂上。比如在一个nn.Sequential里定义了某个层但forward没有调用它钩子自然不触发。再比如模型被torch.compile或JIT化后某些钩子行为会改变这在PyTorch 2.x版本下尤其要注意。还有一个经常被忽视的点钩子对象需要被保存。如果你的register_*_hook()返回值没有被任何变量引用Python的垃圾回收可能把它回收掉钩子随着失效。在交互式环境Jupyter里更容易遇到这种情况。所以一定要写handle net.conv1.register_forward_hook(...)并且在不需要的时候调用handle.remove()。3. 进阶实战一梯度修改、裁剪与回填的标准套路3.1 场景只想对某一层做梯度裁剪torch.nn.utils.clip_grad_norm_是全局梯度裁剪但有时你只想裁剪某个子模块的梯度。比如GAN训练中生成器的某个残差块梯度异常而判别器正常又比如混合精度训练中某些层梯度爆炸其它层没问题。全局裁剪会拖累所有层的学习速度局部裁剪更精准。用Tensor hook实现局部裁剪非常干净import torch import torch.nn as nn from torch.nn.utils import clip_grad_norm_ def make_clip_hook(max_norm): def hook(grad): return clip_grad_norm_([grad], max_norm)[0] # 返回裁剪后的梯度 return hook model SomeModel() handle model.specific_layer.weight.register_hook(make_clip_hook(1.0))注意clip_grad_norm_接收一个参数列表并返回裁剪后的总范数。我在这里用了[grad]把当前梯度包成列表原地裁剪后取回第一个元素作为新梯度返回。这是一个“少有人写但确实可行”的技巧。更稳妥的写法是手动计算def clip_hook(grad): max_norm 1.0 total_norm grad.norm().item() if total_norm max_norm: return grad * (max_norm / (total_norm 1e-6)) return grad这个手动版本的好处是你能在打印日志时看到总范数方便确认hook确实在工作。3.2 场景梯度衰减与正负号干预有些损失函数设计想让某个分支的梯度以一定比例衰减比如让多任务学习中的某个任务“慢一点学”。最笨的方法是在loss里加权重但这样会同时影响该任务所有相关层的梯度。如果只想控制从某个特定张量往后传的梯度大小那Tensor hook是最精确的位置。举一个例子模拟项目Y中我让一个共享encoder的特征向量经过两个head分别做分类和回归。回归任务常常因为样本量少而学得太快影响共享特征。我就在回归head的输入特征上挂了一个缩放hookdef scale_grad_by(grad, factor0.1): return grad * factor feat shared_encoder(x) handle feat.register_hook(lambda grad: grad * 0.1) logits_reg reg_head(feat)这样分类任务回传的梯度不受影响回归任务对共享encoder的贡献被主动削弱。其实这与在loss中乘系数不同后者会把回归loss的所有反向传播梯度同步放大或缩小但不会改变回归head内部参数的相对梯度方向而hook挂在共享输入上只影响“经过这个张量的梯度”有更精确的控制边界。3.3 场景处理retain_graph都解决不了的梯度释放问题默认情况下PyTorch在backward()后会释放计算图对于非leaf Tensor的梯度会被丢弃方便下一次迭代。如果你在训练循环里做了两次backward比如GAN要先更新判别器再更新生成器很多时候需要设置retain_graphTrue这时内存压力会比较大。利用hook可以在某些情况下避免重复backward。举个例子如果两个loss共享同一个骨干网络但你想让第一个loss只更新head A第二个loss更新head B和骨干理论上可以用“stop-gradient”类的操作但更灵活的是在第一个loss反向传播前对head A输出后的共享特征张量挂一个将梯度置零的hookfeat backbone(x) feat.register_hook(lambda grad: torch.zeros_like(grad)) loss1 head_A(feat).sum() loss1.backward() # head_A更新backbone的梯度被清零不影响后续其他loss反向传播这个方案有争议因为零梯度仍然会继续向前传播徒增计算量。如果想真正截断应该用feat.detach()或torch.where(True, feat, 0)这类方式。但零梯度hook在你要“保留计算图结构不重新forward”的场景下有独特价值通过hook先清零再backward之后的第二次backward可以复用计算图但又不会重复累加骨干梯度。当然实际使用时需要对自己的梯度累加逻辑非常清楚我建议先用简单模型验证再上大模型。3.4 场景把梯度“回填”到参数上有时你不想在backward()之后遍历named_parameters()处理梯度而是希望梯度在产生的同时就被“改造”成最终形态。比如你想统计梯度的一阶矩和二阶矩或者把梯度量化后再更新这时可以用hook把梯度保存到一个全局字典里供优化器自定义逻辑使用。grad_storage {} def store_grad(name): def hook(grad): grad_storage[name] grad.detach().clone() return grad return hook handles [] for name, param in model.named_parameters(): if param.requires_grad: handles.append(param.register_hook(store_grad(name)))注意这里必须是detach().clone()。如果只写clone()保存的张量仍与计算图有连接多次backward后可能造成内存累积。另一方面如果你想在hook里检查梯度并在异常时“回填”梯度直接修改返回的grad即可。很多人误以为hook只能“看”不能“改”其实返回值就是用于继续传播的梯度拥有完全控制权。4. 进阶实战二特征图与梯度的诊断统计定位训练异常4.1 用hook采集多层的梯度范数定位NaN源头训练遇到NaN时最痛苦的是不知道NaN是从哪一层先冒出来的。一个高效的排查方案是在每一个卷积层注册full_backward_hook计算每个模块输出梯度的范数和均值实时打印。下面是我常用的诊断代码骨架import torch import torch.nn as nn def make_grad_monitor(module_name, log_freq50): def backward_hook(module, grad_input, grad_output): if grad_output is None: return grad_t grad_output[0] if grad_t is None: return fn torch.isnan(grad_t).float().mean().item() norm grad_t.norm().item() if fn 0 or norm 1e10: print(f[nan/explode] {module_name}: nan_ratio{fn:.6f}, norm{norm:.6g}) return backward_hook for name, module in model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear, nn.BatchNorm2d)): module.register_full_backward_hook(make_grad_monitor(name))大多数情况下NaN并不是所有层同一时刻一起爆掉的而是从某一层开始梯度变成NaN然后顺着反向传播传播到前面的层。通过这种逐层监控你能看到第一个报警的模块名那通常就是问题根源。我在模拟项目X里曾经用这种方式定位到是transformer里某个attention权重计算溢出了——QK之后没有缩放数值直接跑到几千softmax之后的梯度出现了Inf。4.2 用forward hook提取特征图做中间激活的可视化很多同学以为提取特征图必须在网络内部保存或者要重新forward一遍。其实用forward hook完全可以在模型forward完成之后从外部拿到任意层的输出。activation {} def save_activation(name): def hook(module, input, output): activation[name] output.detach() return hook model.conv1.register_forward_hook(save_activation(conv1)) model.layer2.register_forward_hook(save_activation(layer2)) _ model(x) feat activation[conv1] # (B, C, H, W)这里有几个工程化小技巧一定要detach()否则activation字典会持有整张计算图内存可能爆炸。如果只想存最后一次forward的结果记得在每次forward前activation.clear()。如果在一批样本上提取特征做聚类最好不要把特征直接存到增量的list里而是提前分配一个大Tensor用索引写入避免反复append带来的显存碎片。拿到特征图后最常见的处理是对Channel维做归一化再保存成图片import torch.nn.functional as F feat_min feat.min() feat_max feat.max() feat_norm (feat - feat_min) / (feat_max - feat_min 1e-8) grid torchvision.utils.make_grid(feat_norm[0].unsqueeze(0), nrow8) torchvision.utils.save_image(grid, feat.png)4.3 用hook统计梯度分布而不是只看范数范数只能告诉你“大还是小”不能告诉你“梯度里有没有异常模式”。我更推荐在关键层统计梯度直方图的分位数比如打印p50/p95/p99。因为有时候范数正常但99%的维度已经消失只有1%的维度在传梯度这种情况对训练是灾难。用hook实现分位数统计很简单def quantile_hook(module_name): def hook(grad): flat grad.flatten().abs() if flat.numel() 0: q torch.quantile(flat, torch.tensor([0.5, 0.95, 0.99]).to(grad.device)) print(f{module_name}: p50{q[0]:.4g}, p95{q[1]:.4g}, p99{q[2]:.4g}) return grad return hook把这个hook挂到容易出现梯度消失的层训练前期观察数值变化能第一时间发现初始化是否合适。比如一个Sigmoid激活块的p99长期小于1e-8基本可以断定这一块进入饱和区考虑换激活函数或加残差。4.4 自动化把hook采集的数据汇总成训练日志Hook的代码容易散落在训练脚本各处我建议封装成一个“诊断器”类集中管理handle和采集结果class GradientInspector: def __init__(self, model, module_types(nn.Conv2d, nn.Linear)): self.handles [] self.records {} for name, module in model.named_modules(): if isinstance(module, module_types): self.handles.append( module.register_full_backward_hook(self._make_hook(name)) ) def _make_hook(self, name): def hook(module, grad_input, grad_output): g grad_output[0] if g is not None: self.records[name] { norm: g.norm().item(), mean: g.mean().item(), absmax: g.abs().max().item(), } return hook def get_report(self): return self.records def clear(self): self.records.clear() def remove(self): for h in self.handles: h.remove()在训练循环的某个频率下调用get_report()记录到TensorBoard或wandb基本就是一条轻量级的梯度监控流水线。不需要再侵入模型内部的forward代码。这也是“进阶篇”该有的样子不只写得出hook还要把它变成可复用的工具。5. 我最想提醒的Hook坑inplace、作用域顺序与remove机制5.1 inplace操作会让hook的梯度指向被覆盖的张量我在做模型压缩实验时为了省显存在某个层的输出后面直接用了relu_(x)这类inplace操作。结果发现hook拿到的输出和张量本身总是对不上某些统计结果看起来“差一点点”。问题根源是inplace操作会直接在原内存上修改张量而计算图的Function节点记录的仍然是旧Tensor对象引用。反向传播时hook从某个节点拿到的梯度可能对应的是修改前还是修改后的版本行为取决于具体实现完全不可预测。尽量在hook涉及的路径上避免inplace操作。比如ReLU用nn.ReLU()而不是torch.relu_(x)在Transformer的残差连接处不要用x attn(x)改为x x attn(x)。现代PyTorch对inplace操作支持得越来越好但在hook这种“旁路读取”场景中inplace会带来诡异的别名问题排查成本非常高。如果为了显存必须用inplace那就要保证hook只用来观察不要用来修改梯度。5.2 多个hook的顺序不是你想的那样在同一个Tensor或Module上挂了多个hook触发顺序如何PyTorch的文档并没有给非常严格的全局顺序承诺实际经验是多个tensor.register_hook会按注册顺序执行前一个返回的梯度作为后一个的输入Module多个forward hook也会按注册顺序依次执行。但你不能因此就写“依赖顺序”的业务逻辑官方没有把顺序当作公开契约。跨Tensor和跨Module的hook顺序就更难预测了因为反向传播本身是图的拓扑序不同分支的先后并不固定。这个坑在实际中非常常见。比如你想先缩放梯度再做梯度裁剪如果你把它们挂成两个hook依赖执行顺序是不稳定的。稳妥做法是只写一个hook内部自己做两步操作def combined_hook(grad): grad grad * 0.1 grad clip_by_norm(grad) return grad5.3 hook在nn.Sequential和torch.compile下的行为差异nn.Sequential本身是一个Module但hook你通常想挂在里面的子模块上这没问题。真正麻烦的是PyTorch 2.x推出torch.compile后模型会被重写或融合很多子模块在运行时并不以Python对象形式存在此时register_forward_hook可能不触发因为前向已经被编译成融合的C算子。你在调试一个compile过的模型时如果某个hook静默失效不要怀疑是自己写错了先试试在非compile模式下运行。我的经验是训练和推理用compile加速调试和Hook分析关掉compile。两者分开既保速度又保可观测性。如果你必须要在compile模式下挂hook可以试试用torch._dynamo.config.export_mode或者对指定模块加torch.compiler.disable装饰让那个模块用原始Python执行但这是比较偏门的做法新手阶段不建议深挖。5.4 handle.remove()与生命周期管理我在2.4节提过handle要保存这里再展开说。很多人写论文代码时在训练循环里注册hook循环结束后不remove然后反复实例化模型做多组实验结果发现hook越来越多内存越占越高梯度统计也越来越乱。原因就是每次注册产生的handle被丢弃但注册过的hook不会被垃圾回收自动清理除非对应的模型对象被销毁。建议写一个统一的hook管理类所有hook的handle都放进去训练循环结束时统一调用remove()。比如hook_handles [] def add_hook(module, hook_fn): h module.register_forward_hook(hook_fn) hook_handles.append(h) def clear_hooks(): for h in hook_handles: h.remove() hook_handles.clear()这个习惯能帮你避免很多莫名其妙的“它为什么影响了另一个实验”的问题。另外不要在临时函数里注册hook却不返回handle比如def setup(): model.conv1.register_forward_hook(fn) # 不保存handle危险这种写法在交互式环境里经常让hook残留影响后续所有输入的前向输出。5.5detach、requires_grad与hook返回值的关系最后一个容易踩坑的点是即使你在hook里返回了新的梯度原Tensor的.grad也只会在它本身是leaf Tensor时被保存。对非leaf Tensor来说hook处理后返回的梯度会继续向前传播但不会永久保存在该Tensor的.grad属性里。很多人在中间层张量上挂hook期望返回梯度后这个张量有.grad可以读结果发现是None误以为hook没生效。其实hook生效了只是非leaf Tensor默认不保留梯度。如果你希望在hook里既修改梯度又保存梯度必须像前面那样grad.detach().clone()自己存到外部数据结构里。不要试图通过设置tensor.retain_grad()来同时兼得——和hook一起使用时要先明确retain_grad()的作用是让非leaf Tensor保存反向梯度它和hook的返回修改不冲突但会让显存占用上升。6. 收尾一个从调试到落地的Hook应用模板最后分享一个我自己的做法。我通常会把Hook代码独立成一个模块不写在训练脚本里方便后续复用到其它项目。整个“调试型Hook”大概分三层第一层是“记录层”只收集不修改。比如forward hook提取特征图、backward hook统计梯度分位数。 第二层是“干预层”根据记录层的统计结果动态决定是否修改梯度。比如检测到梯度范数超过阈值就放缩梯度。 第三层是“清理层”训练结束后统一remove所有handle输出统计报告。把这个模板写到你的项目里会少走很多弯路class DebugHookManager: def __init__(self, model): self.model model self.handles [] self.storage {} def register_forward(self, name, module): def hook(module, input, output): self.storage[fout:{name}] output.detach().clone() self.handles.append(module.register_forward_hook(hook)) def register_backward(self, name, module): def hook(module, grad_input, grad_output): g grad_output[0] if g is not None: self.storage[fgrad:{name}] g.detach().clone() self.handles.append(module.register_full_backward_hook(hook)) def remove(self): for h in self.handles: h.remove() self.handles.clear()当然这个模板只是一个起点。Hook真正的威力在于它允许你把“观察逻辑”和“模型结构”解耦。以前调试一个几十层的模型我要在每层后面写一堆print和assert改代码时心惊胆战生怕漏删。现在写好在外部模型完全不用动。训练异常时挂上去定位完移除干净利落。如果你刚接触Hook我的建议是从一个最简单的实验开始在一个两层线性模型上挂tensor.register_hook手动算一下梯度验证返回值和原始梯度的关系。之后再过渡到真实模型。不要一上来就在几十层的网络上搞花活因为Hook的调试本身也有成本局部出错很难定位。先把Tensor、Module两套接口的触发时机吃透后面的进阶用法就都是顺水推舟了。