PyTorch 1.8+ 推荐用 register_full_backward_hook 捕获每层梯度,可稳定获取 grad_input 和 grad_output,尤其适合监测输入侧梯度;需判空、命名模块、避免频繁 .item(),优先用 .norm(2).float().item() 统计,并注意 BN/激活层及混合精度下的特殊行为。

用 register_full_backward_hook 捕获每层梯度值
PyTorch 1.8+ 推荐用 register_full_backward_hook 替代已弃用的 register_backward_hook,它能稳定拿到 grad_input 和 grad_output,尤其适合观察某一层输入侧的梯度(即反向传播流经该层时的梯度张量)。注意:它只在该模块的 backward 被调用时触发,且对 torch.nn.Sequential 中的匿名模块需提前命名或遍历获取。
实操建议:
- 钩子函数签名必须是
hook(module, grad_input, grad_output);grad_input是传入该层的梯度(即上层输出的梯度),通常用来判断是否消失 - 避免在钩子内做重计算(如
.mean().item()频繁调用),否则拖慢训练;改用.norm(2).item()或统计非零元素比例更鲁棒 - 对
Linear、Conv2d等参数化层,grad_input[0]对应输入特征图的梯度,grad_input[1]是权重梯度,grad_input[2]是偏置梯度(若存在)
记录各层梯度 norm 并识别消失点
梯度消失不是“全为零”,而是范数持续远低于前几层(比如连续 3 层 grad_input[0].norm(2)
示例关键片段:
立即学习“Python免费学习笔记(深入)”;
grad_norms = {}
def hook_fn(module, grad_input, grad_output):
if len(grad_input) > 0 and grad_input[0] is not None:
# 只取输入特征图的梯度(跳过权重/偏置梯度)
g = grad_input[0]
grad_norms[module._get_name()] = g.norm(2).item()
<p>for name, module in model.named_modules():
if len(list(module.children())) == 0: # 叶子模块
module.register_full_backward_hook(hook_fn)</p>注意:grad_input 可能含 None(如某输入无梯度),务必判空;module._get_name() 比 name 更可靠,因 name 在 Sequential 中可能是空字符串。
为什么不用 torch.autograd.grad 手动求梯度
手动调用 torch.autograd.grad 对每层输出求关于损失的梯度,理论上可行,但实际会破坏计算图、引发 RuntimeError: Trying to backward through the graph a second time,除非显式设 retain_graph=True —— 这会导致显存爆炸,尤其在深层网络中不可行。
更严重的是,它无法反映真实训练时梯度流动路径:例如某层被 torch.no_grad() 包裹、或中间有 .detach(),手动求梯度会绕过这些控制,给出错误“有梯度”假象。而 register_full_backward_hook 完全复现真实反向流程,结果可信。
BN 层和激活层的梯度陷阱
BatchNorm 的 running_mean/running_var 不参与反向传播,但其 weight 和 bias 有梯度;而 ReLU 类激活层本身无参数,grad_input[0] 就是上游梯度乘以导数掩码(负值处为 0),所以它的 grad_input[0].norm() 天然比前一层小——不能单看绝对值,要对比同位置通道的衰减率。
容易踩的坑:
-
Sigmoid/Tanh后接线性层时,梯度常在前几轮就坍缩到 1e-10 以下;此时看grad_output[0](该层输出的梯度)比grad_input[0]更能定位问题源头 - BN 层在
eval()模式下不更新统计量,但梯度仍正常流动;误以为“BN 关闭导致无梯度”是常见误解 - 混合精度训练(
amp)下,grad_input是 FP16 张量,.norm(2)易下溢,建议先转.float()再算范数
真正难排查的是跨层耦合问题:比如某 Conv 层权重初始化过小 + 后接 Sigmoid,梯度在第 3 层就归零,但钩子只显示“这一层没梯度”,得结合前序层的 grad_output 和本层输入分布一起看。


















