资讯动态

计算图与反向传播:从梯度消失到训练稳定性的排查指南

发布时间:2026/10/19 10:20:27 来源:尧图企业网站定制
前阵子有个朋友在群里发了一张loss曲线问我为什么训练到一半梯度突然变成了NaN。我让他把backward之后的每个参数梯度打印出来看一眼果然好几层的梯度全是NaN。这种事在深度学习调试里太常见了而真正能让你稳住不乱去排查的底层知识恰恰就是题图里那两个字计算图和反向传播。这两样东西听起来像教材里的基础概念可一旦实际训练出问题不懂的人只能靠反复改学习率碰运气懂的人直接沿着图往回追一遍就能定位。这篇文章不是教科书式的原理堆砌而是想用“干活的人”的视角把计算图和反向传播彻底讲透——从它是什么、怎么运作到手推一遍梯度再用数值梯度做一个验证工具最后聊聊我在真实训练里踩过的梯度的坑以及一套用了很久的排查习惯。适合刚入门深度学习、对框架里那行backward()半懂不懂的读者也适合有段时间没碰原理、正在被各种训练事故折腾的工程师。哪怕你只写过几行模型也能照着文章里的方法亲手验一遍梯度发现自己对模型的掌控感完全不一样。1. 计算图的本质前向计算如何被拆成可追踪的积木1.1 计算图到底在记录什么我先直接用一句话定义计算图是一张有向无环图把一次完整的前向计算拆成“节点”和“边”。节点是运算边是节点之间流动的张量数据。拿最常见的一个小表达式来感受一下y x * w b拆开就是三个节点x * w是一个乘法节点结果记为tt b是加法节点结果记为y。x、w、b 是叶子也就是没有父节点、不依赖其他计算的输入。整个前向计算就是从叶子出发沿着依赖方向一路推到根节点y。这套抽象看起来不复杂但它解决了一个特别关键的问题神经网络本质是巨型嵌套函数几百万个参数叠在一起如果依赖“人类手推公式”去解梯度三层以上的网络就已经痛苦到劝退了。但把计算过程拆成一个一个小节点之后每个节点都只负责一个极其简单的局部运算比如相乘、相加、取指数、做除法。这些简单运算的导数是绝对可算的于是复杂函数的求导就变成了在这些简单节点的导数之间做链式传递。所以计算图的第一个价值是简化求导第二个价值是便于复用。同一个节点可以被多个上游节点引用比如某个特征同时参与了两个 loss 分量的计算图结构会自动把这个事实记录下来反向传播时梯度会自然地在分支处相加。这些逻辑如果用裸代码来实现很容易漏而图结构天然能兜住这种复杂度。1.2 前向过程就是一次有趣的图遍历你可以把前向传播理解成“按拓扑序遍历这张图”从叶子出发每经过一个节点就算一次局部结果把结果往下游传直到根节点输出 loss。重要的是遍历的同时框架会顺手把两样东西挂到张量上当前节点的输出值以及这个输出对应的运算信息。在 PyTorch 里运算信息就是每个张量头上的grad_fn属性。我打个比方计算图像一张菜谱步骤图。输入是原始食材每一步操作是切菜、炒菜、调味成品是最终的 loss。如果菜出锅味道不对你不可能把整桌菜倒回去重新找原因但你可以沿着步骤往回查是哪一步盐放多了还是哪一步火候过了。计算图给深度学习提供的正是这种“回溯路径”。你不仅知道最后 loss 是多少还知道它是经过哪些步骤算出来的每一条路径都清晰可查。这里有个初学者很容易忽略的细节并不是所有张量都需要挂grad_fn。只有当你设置了requires_gradTrue并且这个节点参与运算后产生的新的张量才会被记录进计算图里。普通数据的读取、转置、类型转换如果不影响梯度链就不会额外产生图节点。这也是为什么我们常看到某些中间变量有grad_fn某些没有——没有的要么本来就是叶子要么已经被框架判定为“不需要跟踪”。1.3 动态图和静态图两种建图哲学的取舍我早期写过静态图框架也长期用动态图框架这两者的差别对理解计算图特别有帮助。静态图框架的思路是你先描述好整个计算流程框架一次性把整张图建好、编译优化之后真正执行的是这张编译后的图。好处是它可以做全局优化比如算子融合、内存规划执行效率高坏处是调试起来很别扭你没法在中间步骤随便打印一个临时值因为图是提前被定义好的。动态图框架则相反每跑一次前向就现场把图搭建一遍。你的 Python 代码怎么写图就怎么长。这对调试和实验极其友好可以打印任意中间值可以用if语句控制分支甚至可以临时插入一个print而不影响整体。代价是每次迭代都有建图的开销极端情况下会成为训练的瓶颈。用一个生活化的类比静态图像先画好施工图纸再施工图纸画完就定死动态图像即兴演奏边听边调灵活但每次都要重新准备乐器。我个人的建议是日常研究和调试用动态图省时间真正到了部署和上线阶段再把训练好的计算流程出口到支持静态推理的引擎上两边的好处都吃到。理解这个区别有助于你在不同场景下选择合适的技术路线。2. 反向传播的数学直觉从损失函数倒推参数梯度2.1 反向传播的唯一数学工具链式法则反向传播的全部秘密其实就一条复合函数求导的链式法则。如果你面前有一个嵌套函数L f(g(h(x)))那么它对 x 的导数就是逐层相乘dL/dx dL/df * df/dg * dg/dh * dh/dx反向传播之所以叫“反向”是因为它从最外层的 loss 出发逐层向内回传。每一层在回传时做的事情只有两件先从上游拿到“loss 对这个节点输出的梯度”再计算“这个节点针对它自身输入的局部梯度”两者相乘就得到了送往更下游的梯度信号。这像一个多级减速传动系统最外端的输出转速变化会通过每一级齿轮的传动比一路传回输入端。总传动比等于每一级传动比的乘积。计算图里每一段“传动比”就是局部导数反向传播就是一次性把所有传动比从后往前连乘出来。很多人在这一步卡住是因为一下子看到太多符号不知道乘法顺序怎么安排。我的建议是只记住一句话上游梯度永远乘以本地的导数。上游梯度是框架传下来的、已经算好的本地导数就是你当前节点输出对自己输入的偏导。至于方向你只需要判断“当前操作的输出是什么输入是什么”然后把对应的导数找出来乘上去其他都不用想。2.2 手推一个双层全连接从最外层一路剥到第一层公式看得再多不如亲手推一层真实的网络。我下面推一个最基础的双层全连接网络你用笔照着纸上走一遍感受会完全不一样。网络定义如下输入X形状[N, D]第一层参数W1形状[D, H]、b1形状[H]激活函数用 ReLU第二层参数W2形状[H, C]、b2形状[C]输出logitsloss 用交叉熵或者 MSE 都行。前向过程h relu(X W1 b1) logits h W2 b2 L loss(logits, target)现在从后往前推梯度。先看最后一个线性层logits h W2 b2。如果已知dL/dlogits那么dL/dW2 h^T dL/dlogits dL/db2 sum(dL/dlogits, axis0) dL/dh dL/dlogits W2^T注意中间这一步dL/dh的求法是在原来的矩阵乘法方向上做转置操作h W2变成dL/dlogits W2^T形状刚好是[N, H]与 h 一致。接着要把梯度穿过 ReLU 节点。ReLU 的导数是分段函数输入大于 0 时导数为 1否则为 0。于是dL/dz dL/dh * (z 0)这里的z X W1 b1(z 0)是一个和 z 形状完全相同的 0/1 掩码乘法是按元素乘。最后回到第一层参数dL/dW1 X^T dL/dz dL/db1 sum(dL/dz, axis0)你看整个过程就是不断重复“找局部导数乘上游梯度再往上游传”。一旦形成一个清晰的模板哪怕后面换成三层四层、加上各种残差连接本质上也还是同一套逻辑在重复。2.3 梯度形状必须对得上参数形状手推的时候大部分人第一次犯错不是计算出错而是形状对不上。梯度本身也是一个张量它的形状必须和对应参数完全一致否则优化器 step 的时候就会报错或者更糟——不报错但结果完全不对。我自己习惯用“形状对齐法”来检查公式先确认目标梯度的形状比如dL/dW2一定和W2一致也就是[H, C]。然后再看手头有哪些已知形状h是[N, H]dL/dlogits是[N, C]。要得到[H, C]唯一的乘法方式就是h^T dL/dlogits。这样就不会把乘法的左右顺序搞反。还有一个容易混淆的点对偏置项的梯度通常需要对 batch 维度求和因为偏置是广播到每一个样本上的。例如b2的形状是[C]而dL/dlogits是[N, C]那么dL/db2应该是对 N 维求和得到[C]。很多自动求导框架会直接帮你把这一步也算好但如果需要手写某个自定义算子这个小细节就很可能在工作里让你多花一个小时去排查。我建议所有刚开始研究反向传播的读者都至少亲手推一次完整的两层网络再用框架验证每个参数的梯度值。这个过程大概一小时收益却是长期的。你会开始真正信任框架的自动求导不再把它当黑盒也会在处理自定义算子、梯度修改、特殊结构的时候心里有底。3. 动手验证手写微型自动微分并用数值梯度校验3.1 最简验证标量计算图手算加跑代码理论推完我用一个最小的例子来验证。假设你有一组标量输入计算关系是import torch x torch.tensor(2.0, requires_gradTrue) w torch.tensor(3.0, requires_gradTrue) b torch.tensor(1.0, requires_gradTrue) y x * w b loss (y - 4) ** 2 loss.backward() print(x.grad , x.grad) # 期望输出 18.0 print(w.grad , w.grad) # 期望输出 12.0 print(b.grad , b.grad) # 期望输出 6.0我手算一遍给你看。先算前向y 2×31 7loss (7-4)^2 9。再算反向dL/dy 2×(7-4) 6。接着dy/dx w 3所以x.grad 6×3 18dy/dw x 2所以w.grad 6×2 12dy/db 1所以b.grad 6。和代码输出完全对得上。这段代码虽然短但它把整个计算图的运作讲明白了每个变量后面挂着 grad反向传播通过链式法则把 loss 对每个叶子节点的梯度都算了出来。你可以在这一步停顿一下把计算图画出来亲手走一遍。不要在纸上想直接动手写一次记忆效果完全不同。3.2 数值梯度给自动求导做的尺子自动求导会不会算错理论上不会但当你写了自定义算子、修改了前向规则或者怀疑某个地方莫名其妙把梯度挡住了就需要一把独立的尺子来验证——这就是数值梯度。数值梯度基于导数的定义用很小的步长eps分别计算出f(xeps)和f(x-eps)两者之差除以2*eps这就是中心差分近似。我在实际验证中用的就是这个方式因为比单侧差分更稳定。下面给一个最小实现def numerical_grad_of(loss_fn, param, index, eps1e-5): flat param.detach().clone().flatten() flat[index] eps loss_plus loss_fn(flat.reshape_as(param)).item() flat[index] - 2 * eps loss_minus loss_fn(flat.reshape_as(param)).item() flat[index] eps return (loss_plus - loss_minus) / (2 * eps)使用上你可以对一个简单模型挑某个参数的一个元素来验证把该元素做微小扰动重新算 loss得到数值梯度再调用自动求导得到解析梯度两者做差如果相对误差小于1e-4左右说明这条梯度路径是健康的。我在写自定义torch.autograd.Function时几乎必跑一次这套检查而且在训练不收敛的时候也会拿它对可疑层做一次体检。它就像是做菜时旁边放的一个测盐度的小试纸不需要特别精确但能在你怀疑“盐加多了”的时候立刻给你一个量化答案。3.3 梯度审计的实操流程与适用范围实际工程里我不建议你对整个大模型每个参数逐个元素做数值梯度那会慢到怀疑人生。我的日常做法是在模型还是玩具规模的时候或者新引入某个自定义模块的时候做一次完整审计。代码思路大概是这样def audit(model, data, target, eps1e-5): model.train() logits model(data) loss F.cross_entropy(logits, target) loss.backward() for name, param in model.named_parameters(): if param.grad is None: print(f{name}: grad is None) continue analytic param.grad.detach().clone() flat param.detach().clone().flatten() numeric torch.zeros_like(analytic) for i in range(flat.numel()): flat[i] eps loss_plus F.cross_entropy(model(flat.reshape_as(param)), target).item() flat[i] - 2 * eps loss_minus F.cross_entropy(model(flat.reshape_as(param)), target).item() flat[i] eps numeric.flatten()[i] (loss_plus - loss_minus) / (2 * eps) diff (analytic - numeric).abs().max().item() print(f{name}: max_abs_diff {diff:.6e})这个脚本只适合小模型几百个参数以内跑得动。我的经验是不要指望它在生产模型上跑完而是把目标缩小到单独一层或者一个自定义算子上。框架自带的torch.autograd.gradcheck本质上也是在做类似的事只是接口更严谨把某些容易出错的边界情况也包含了。如果你要写自定义Function直接用gradcheck是更省力的方案。把数值梯度作为验证手段之后你会获得一个珍贵的技能可以直接判断一个反向传播实现是否出错而不是靠猜。这种事在真实项目里经常遇到尤其是当某个新同事实现了一个看起来很 fancy 的模块却总是训练不收敛的时候。这时候我一般不慌着调超参先跑一遍梯度审计通常五分钟之内就能确定是不是梯度本身写错了。4. 训练闭环梯度如何被优化器变成参数更新4.1 一次训练迭代的五步闭环理解了计算图和反向传播再回头看训练循环会特别清晰。常规的一个 iteration 分为五步logits model(data) loss loss_fn(logits, target) optimizer.zero_grad() loss.backward() optimizer.step()前两步是前向第三步清空旧梯度第四步反向传播第五步参数更新。初学者最常犯的错误之一就是忘记zero_grad。backward()默认会把梯度累加到.grad属性上这是有意设计的用来支持梯度累积这种特殊用法。但如果你每次迭代都不清零那这个累加行为就会让梯度方向里混进之前好几步的信息loss 曲线就会呈现一种混乱的持续上涨或者剧烈震荡状态。我实测过一个很典型的场景一个简单的 MNIST 分类模型某次我删掉了zero_grad训练 loss 不降反升五分钟内模型参数就飞了。之后我再也没忘过这一步。4.2 优化器是怎么消费梯度的很多人以为优化器 step 就是简单减去lr * grad其实不同优化器对梯度的“加工方式”差别很大。SGD 就是最朴素的用法θ θ - lr × g梯度方向直接决定更新方向。动量方法则在梯度方向上叠加了历史速度v_t β * v_(t-1) g θ θ - lr * v_t到了 Adam 这一代它更进一步分别计算梯度的一阶矩均值和二阶矩平方均值然后用矫正后的值做归一化更新。换句话说Adam 不再关心梯度的绝对大小而更关注梯度的相对方向所以它对学习率的选择比 SGD 更宽容一些并且天然对梯度量纲变化更鲁棒。不管优化器内部怎么做它们的输入都是反向传播产出的梯度。所以你可以这样理解反向传播决定方向优化器决定步长和路径的平滑程度。如果某一个参数的梯度方向是错的哪怕是 Adam 也无法挽回。我在调试时有一个习惯在loss.backward()和optimizer.step()之间临时插入一段打印看看当前每个参数的梯度均值和标准差。这能帮助我区分“梯度方向混乱”还是“梯度量级异常”前者可能是网络结构或者 loss 函数有问题后者更可能和学习率、数据分布有关。4.3 梯度范数是一切训练稳定性的晴雨表训练神经网络我最先看的一个监控指标不是 loss而是全局梯度范数。原因很简单loss 是一个标量它不告诉你问题是出在局部还是出在整体而梯度范数刻画了“当前参数空间被推动的力度”到底有多大。如果梯度范数突然从正常水平跳到1e4几乎是铁定的信号——某个层出现了数值爆炸。这时再去看 loss通常已经 NaN 或者震荡了。计算梯度范数的代码很简单total_norm 0.0 for p in model.parameters(): if p.grad is not None: param_norm p.grad.detach().data.norm(2) total_norm param_norm.item() ** 2 total_norm total_norm ** 0.5 print(total_norm)PyTorch 里也有现成的函数torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)它除了计算范数还能在范数过大时直接裁剪梯度。这个方法可以在梯度爆炸的早期把训练稳住。但我得提醒一句clip_grad_norm 治标不治本它只是把瞬时过大的波峰切掉。真正要解决问题还是得靠调低学习率、改进初始化、加归一化层或者检查数据里是否混入了异常值。我建议在每次实验开始时都打印一批正常的梯度范数作为基准值。比如某个模型正常时梯度范数在0.1~1之间训练十轮之后如果突然变成1e5你立刻就知道出了问题而不需要等 loss 彻底崩掉才反应过来。5. 反向传播事故现场常见坑位与排查实录5.1 常见异常症状速查表这一节是我在实际训练里最常用到的东西。我把常见症状、可能原因和排查方向整理成一张速查表症状常见原因排查方向某个参数梯度一直是 0该参数没有参与 loss 计算被.detach()截断ReLU 输入恒负导致死亡打印param.grad检查计算图路径所有参数梯度都是 NaN学习率过大数据里有 NaNloss 函数数值不稳定检查输入数据、降低学习率、改用 log_softmaxloss 不降zero_grad漏写优化器参数列表和模型参数不匹配部分层被锁定先跑通一轮小数据观察梯度范数梯度检查误差很大数值微分步长不合适模型存在不可导点float32 精度不足调整 eps 到 1e-5检查激活函数是否可导早期正常后面梯度异常变大数据分布偏移学习率过大且发生震荡检查输入数据分桶统计启用梯度裁剪训练损失震荡极其剧烈梯度累积未清空batch 大小过小学习率偏高检查每一步是否有 zero_grad调小学习率或增大 batch这张表我贴在工位旁边贴了很久每次排查训练问题先从里面选症状再按对应方向走基本能省掉一半的瞎折腾时间。5.2 几个让我抓狂过的高频坑第一个坑.detach()用错了地方。它本意是切断梯度从当前节点继续回传如果代码作者对这个机制理解不深很容易在网络中间某个不小心的地方插一个detach导致后面的所有参数都收不到梯度。你会看到 model 前面几层参数不变loss 却正常下降相当迷惑。排查方式就是拿第 3 章的梯度审计或者直接打印各层param.grad是否为空。第二个坑inplace 操作破坏计算图。比如x 1或者relu_这种带下划线的原地操作它们会改写参与运算的张量本身而计算图里还保留着旧值的引用backward 时读到已经被覆盖的数据经常报错或者得到错误梯度。框架一般会给出类似“a leaf Variable that requires grad is being used in an in-place operation”的警告。出现了不要硬扛改成x x 1重新赋值或者用clone()后再做原地操作。第三个坑ReLU 神经元“死亡”。当某层权重被更新到偏负区域ReLU 输出恒为 0梯度过不去这个通道就彻底死掉。表面症状是该层输出的分布里出现大量常数 0且梯度为 0。解决手段包括换用 LeakyReLU、调整初始化、降低学习率或者检查是不是权重更新被某次极大梯度带偏了。我排查这类问题的方法是打印激活输出的均值与 0 的比例而不是看 loss 曲线——因为一个通道死了loss 可能还在缓慢下降非常隐蔽。第四个坑忘记在no_grad下操作参数导致计算图越建越长。严格来说这不算反向传播算错但会导致显存占用持续上升、训练越来越慢。这通常发生在你试图在训练中途直接改参数时。正确做法是把参数修改包在torch.no_grad()里或者通过优化器统一管理避免在“需要梯度”的上下文里给参数做赋值。5.3 我平时的梯度调试清单最后分享一下我现在的固定调试流程。每次新建一个实验训练逻辑写完、跑第一遍代码之前我都会按以下清单过一遍用一个极小的 batch比如 2 条数据跑一次完整的前向和反向确认每个参数都有梯度没有 None对可疑的自定义算子或者最近改动的层跑一次数值梯度校验打印一个初始化状态下所有参数梯度的均值、标准差和全局梯度范数存成基线训练前几轮时观察梯度范数和 loss 的联动关系如果范数在正常范围但 loss 不降再去查数据是否正确把zero_grad写成循环的第一步而不是放在backward后面避免漏掉全程保留torch.autograd.set_detect_anomaly(True)的习惯仅在调试时让框架在反向传播出现 NaN 时第一时间报出错的位置。这套清单帮我解决过至少十几次看起来毫无头绪的训练事故也让我在给别人排查模型问题时有了更结构化的思路。很多看起来玄乎的神经网络训练问题追到最后都发现不是模型结构的问题而是反向传播这条链路上某个不起眼的位置断了、错了、或者形状没对上。把计算图和反向传播当成排查工具去掌握比背一万个网络结构都更管用。我自己的体会是对计算图和反向传播的理解深度直接决定了你在深度学习调试时候的从容程度。与其遇到问题再临时翻文档不如现在就照着文里的例子手推一遍再用数值梯度验证一次。这套基本功打扎实之后你再看那些五花八门的训练技巧会发现它们大都是围绕“如何让梯度更稳、更准、更有效地流动”这一个中心问题在转。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑