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)
示例关键片段：
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 和本层输入分布一起看。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
