很多人在初学PyTorch时遇到的第一道坎往往是梯度计算。明明照着文档写好了前向传播调用一次backward()得到的梯度却和自己手推的结果对不上或者报错信息让人摸不着头脑。这篇文章我结合自己实际调模型的经验把PyTorch中梯度计算的完整链路从头到尾拆一遍从requires_grad的底层规则到backward()执行时的向量-雅可比积再讲到梯度累积、参数冻结、detach和自定义反向传播。无论是刚入门想搞懂原理的初学者还是被梯度丢失、显存爆炸折磨的调参选手这篇都应该能帮到你。1. 梯度计算的地基从偏导数到PyTorch的requires_grad1.1 梯度计算为什么不能只靠“背公式”刚接触深度学习的时候我们首先在纸上算导数。比如函数 $y x^2 3x$在 $x2$ 处的导数是 $7$。这个计算很简单但真实模型里损失函数往往由几十层神经网络复合而成参数动辄百万级。你不可能为每一层手动推导梯度表达式更不可能在每次更新参数时都重新手算一遍。梯度计算在PyTorch里并不是一个“求导公式集合”而是一套自动微分机制。你只需要把前向计算写清楚PyTorch会记录每一步运算然后从最终损失出发沿计算路径反向走一遍把梯度逐个传递回每个参数。整个过程叫自动微分落在代码层面就是autograd引擎。初学者最容易犯的错是把“梯度”理解为某个具体数值。实际上梯度是一个向量它由损失函数对各个参数的偏导数组成。参数有多少个梯度向量就有多少个分量。比如一个只有两个参数的模型 $z w_1 x_1 w_2 x_2$损失对参数的梯度是 $(\partial L/\partial w_1, \partial L/\partial w_2)$。PyTorch中的x.grad保存的正是这种梯度张量形状和原参数张量保持一致。1.2 requires_grad与叶子节点的规则要让一个张量参与梯度计算最简单的做法是在创建时设置requires_gradTrueimport torch x torch.tensor(2.0, requires_gradTrue) y x ** 2 3 * x y.backward() print(x.grad) # tensor(7.)这段代码里x被称为叶子张量y是由运算得到的结果张量。y.backward()执行后梯度会从y流回x并累积在x.grad中。这里有个容易被忽略的规则只有requires_gradTrue的张量以及与它发生过运算的结果张量才会被autograd追踪。如果两个普通张量做乘法结果是不需要梯度的一旦其中一方requires_gradTrue计算结果就自动进入计算图反向传播时可以拿到梯度。还有一个规则是叶子节点的梯度只累积不覆盖。多次调用backward()如果计算图允许x.grad会把每次反向传播得到的梯度相加。这也是为什么训练循环里每个batch都要调用optimizer.zero_grad()——不把上一次的梯度清零梯度就会一直累加参数更新方向就会乱掉。1.3 “does not require grad”报错到底是什么意思我见过很多人看到这句报错就懵了RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn。这句话翻译过来是你试图对一个不需要梯度的张量执行反向传播或者你的计算图中某个环节断开了梯度追踪。出现这个报错通常有三个原因。第一你创建的张量没有设置requires_gradTrue而你对它的运算结果直接调用了backward()。第二输入数据是从NumPy数组或者Python列表转过来的默认requires_gradFalse又没有手动开启。第三你对某个中间结果调用了detach()把它从计算图中摘了出去后续基于它的计算就丢失了梯度链路。排查思路很简单在backward()之前打印每个关键张量的requires_grad和grad_fn属性看看是在哪一步断掉的。grad_fn是autograd为运算节点分配的反向函数对象普通张量的grad_fn是None而由运算产生的张量一定带有类似AddBackward0这样的信息。2. autograd的链路计算图、grad_fn和向量-雅可比积2.1 动态图是在前向过程中“顺手”搭起来的PyTorch的计算图是动态图意思是它不像某些框架那样先定义好一张静态图再执行而是每执行一行代码就同步记录一步运算关系。这个记录行为由autograd引擎自动完成你甚至不需要显式告诉它“我要建图”。我把这个机制理解成账本。每一次加、乘、幂运算都会在账本里记一笔AddBackward0代表这次操作是加法MulBackward0代表乘法。当调用backward()时PyTorch就从最终输出这个节点开始沿着账本记录的反向依赖关系把梯度像流水一样倒着往回推。实际使用中grad_fn是最直观的观察窗口。你可以打印一个中间结果张量的grad_fn它会显示该节点是由哪个反向函数生成的。例如x torch.tensor(2.0, requires_gradTrue) a x * 3 # a.grad_fn 是 MulBackward0 b a 1 # b.grad_fn 是 AddBackward0 c b ** 2 # c.grad_fn 是 PowBackward0从c出发通过c.grad_fn.next_functions你能一层层回溯到a再回溯到x。这串链条就是梯度计算走过的那条路和你手推复合函数链式法则时写的中间变量完全对应。2.2 backward() 为什么有时必须传gradient参数这是很多初学者第一个大坑。对y.backward()来说如果y是一个标量一切都很自然但如果y是一个向量直接调用backward()就会报错提示你必须传入一个与y形状相同的gradient参数。原因是autograd的根本机制是向量-雅可比积。假设你有一个中间向量 $y (y_1, y_2, ..., y_n)$它由输入 $x$ 计算而来。反向传播时autograd不直接计算 $y$ 对 $x$ 的完整雅可比矩阵而是计算一个外部梯度向量 $v$ 与雅可比矩阵的乘积。这个 $v$ 就是你传入的gradient参数。大多数情况下我们训练的损失函数都是标量所以backward()不需要传参数。但如果你的网络输出的是一个向量并且你想把梯度回传到更早的层就必须自己提供$v$。最常用的$v$是全1向量表示“我希望输出向量的每个分量对损失的贡献权重都是1”。这就是loss.backward(torch.ones_like(loss))这种写法的来源。2.3 用一个手算例子验证grad_fn链路我们来看一个稍微具体一点的例子。设输入为 $x$前向计算为$$u 3x$$ $$v u 1$$ $$y v^2$$如果 $x2$那么手算梯度是$$\frac{dy}{dx} \frac{dy}{dv} \cdot \frac{dv}{du} \cdot \frac{du}{dx} 2v \cdot 1 \cdot 3 2 \times 7 \times 3 42$$用PyTorch验证x torch.tensor(2.0, requires_gradTrue) u x * 3 v u 1 y v ** 2 y.backward() print(x.grad) # tensor(42.)如果你把u、v、y的grad_fn都打印出来会得到一条清晰的链PowBackward0 - AddBackward0 - MulBackward0。反向传播时autograd就是沿着这条链从$y$的梯度1出发依次乘上每一层的局部导数最终得到$x.grad42$。这个例子虽然简单但它完整展示了链式法则在代码层面的落地方式。3. 梯度计算的日常实操backward、优化器和梯度累积3.1 标量损失最省事向量输出怎么处理训练深度学习模型时最终用于反传的通常是标量损失。比如分类任务里的CrossEntropyLoss它接收模型输出和标签返回一个标量。这种情况直接loss.backward()就行autograd会自动把初始梯度设为1从标量损失开始反传。但有些场景下你处理的不是标量比如在对比学习或者多任务学习中你可能想对某个中间层输出做梯度操作。这时候需要手动构造初始梯度。一个常见的做法是outputs model(batch_data) # 形状为 [batch_size, dim] loss criterion(outputs, targets) # 标量 loss.backward()如果出于某种原因你想对outputs本身调用backward()就得传一个和它形状相同的向量outputs.backward(torch.ones_like(outputs))这个操作的含义是把outputs的每个分量都当作同等重要的目标分别求梯度后再求和。如果你的下游处理对每个分量有不同的权重直接把这个权重向量传给backward()即可。3.2 zero_grad、loss.backward、optimizer.step的顺序为什么不能乱标准的训练循环长这样optimizer.zero_grad() loss criterion(model(data), target) loss.backward() optimizer.step()顺序不能乱原因在于PyTorch的梯度累积机制。梯度是累积在张量的.grad属性中的。optimizer.zero_grad()负责把模型所有参数的.grad清零loss.backward()计算出当前batch的梯度并累加到.grad上optimizer.step()根据.grad和优化器内部的动量状态更新参数。如果把optimizer.step()放在zero_grad()前面相当于用上一次的梯度加上本次的梯度一起更新一次参数会来回震荡。如果把zero_grad()放到backward()之后那这一步的梯度会被立刻清零等于没更新。很多新手刚跑模型时loss曲线奇怪地反复横跳八成就是顺序写错了。3.3 梯度累积想变大batch又不想爆显存的做法显存不够是训练模型时最头疼的问题之一。一个常用的解法是梯度累积把多个小batch的梯度加到一起攒够一定数量再更新一次参数。实现起来很直接accumulation_steps 4 model.zero_grad() for i, (data, target) in enumerate(train_loader): loss criterion(model(data), target) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() model.zero_grad()这里的除以accumulation_steps是必须的。如果不除累积多个batch后再更新等效batch size变大但梯度也会按倍数放大学习率就需要跟着调整。除一下让累积后的梯度量级和单个batch差不多训练超参数可以基本保持不变。要注意的是model.zero_grad()和optimizer.zero_grad()在大多数情况下是等价的但如果模型里有某些参数不经过优化器管理用model.zero_grad()会更稳妥。4. 控制梯度流动requires_gradFalse、detach和no_grad4.1 三种做法的区别和使用场景很多人在需要“让某些参数不参与梯度计算”时会把三种方法混为一谈。先说结论requires_gradFalse彻底关闭某个张量的梯度需求。该张量可以继续参与运算但运算结果不会再追踪它的梯度。detach()返回一个新的张量这个张量和原张量共享数据存储但被从计算图中摘了出来。新张量的requires_gradFalse反向传播不会通过它回传。torch.no_grad()一个上下文管理器包裹的代码块内所有计算都不被autograd追踪也不会产生梯度。三者的区别可以从“是否影响数据本身”来理解。requires_gradFalse和detach()都不改变数据值只是改变追踪状态no_grad()则更像一个全局开关适合用在推理阶段能省下大量内存和计算开销。4.2 冻结预训练模型参数迁移学习时我们经常加载一个预训练模型只希望训练最后几层分类头其余层保持不动。最直观的做法是遍历模型参数全部设置requires_gradFalsefor param in model.parameters(): param.requires_grad False # 只训练最后一层 for param in model.fc.parameters(): param.requires_grad True optimizer torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3)这里有一个小技巧优化器只接收requires_gradTrue的参数。用filter过滤一遍能省掉不少无效计算也避免优化器往冻结参数里塞状态信息。但要注意requires_gradFalse只禁止了梯度计算前向传播时如果这些参数参与运算依然会占用显存。如果数据量很大更推荐用torch.no_grad()包住整个前向过程或者直接把模型切分成“冻结部分”和“训练部分”冻结部分先算好特征存起来。4.3 GAN和对比学习里detach的常见用法detach()最常见的应用场景是生成对抗网络。训练判别器时我们希望生成器不接收从判别器传来的梯度训练生成器时又希望判别器参数不被更新。如果直接用no_grad()包住前向会丢掉必要的梯度信息。正确做法是前向传递时对生成器的输出调用detach()。假设你要训练一个生成对抗网络判别器判断真伪损失反传给生成器fake generator(noise) # 训练判别器时让梯度不要传到生成器 d_loss discriminator_loss(discriminator(fake.detach()), target_fake)这里fake.detach()断开了生成器的梯度流判别器可以正常更新而生成器不会收到任何来自这一支的梯度。训练生成器时则相反需要让梯度通过判别器流回生成器因此不能detach()。对比学习里也有类似操作。很多对比损失中某个正样本对既希望当前样本的表示靠近它又希望它在梯度更新时不要影响到另一个分支。这时对其中一个分支的输出做detach()能让两个编码器的更新节奏解耦。说白了detach()就是你在计算图里主动剪断一根线的工具。5. 小心这些梯度计算陷阱in-place修改、retain_graph和数值稳定5.1 in-place操作与版本计数器in-place操作是指在原张量上直接修改数据比如x.add_(1)、x.mul_(2)、x[0] 3。这类操作的问题在于autograd在反向传播时需要读取前向计算时保存的值如果这些值被就地改掉了梯度计算就会出错。PyTorch用版本计数器来检测这种问题。每个张量都有一个版本号每次进行in-place修改版本号就会增加。反向传播时autograd会检查保存的中间结果版本号是否和当前一致不一致就报错RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.。这个坑最常见于ReLU(inplaceTrue)搭配后续需要反向传播的复杂计算。虽然很多正规模型都在用inplaceTrue那是因为实现时小心处理了数据流。新手阶段建议先全程用非in-place版本比如relu(x)而不用relu_(x)免得排查起来无从下手。如果非用in-place不可切记不要对需要梯度的叶子张量做in-place修改也不要对模型的输入做in-place修改。5.2 多分支共享节点时该用retain_graph还是重新forward正常的backward()调用后计算图的缓冲会被释放以节省内存。如果你对同一个计算图调用两次backward()第二次就会报错说图已经被释放。这时有两个选择一是给backward()传入retain_graphTrue保留计算图允许再次反向传播。但是这么做会显著增加内存占用不能无脑滥用。二是重新执行一遍前向得到新的计算图再反向传播。这个方案需要额外算一次前向但内存开销更可控。需要说明的是这两种做法适用的场景不同。如果你是在调试、观察梯度用retain_graphTrue很方便如果是训练只是在某个损失函数里同时用了多个输出分支最好把前向封装成函数多次调用每次都让自动微分创建独立的新图。千万别在训练循环里对所有backward()都加上retain_graphTrue否则显存会迅速膨胀。5.3 梯度消失和梯度爆炸的快速定位方法梯度消失和梯度爆炸是老生常谈的问题。遇到训练不收敛、loss变成NaN、模型参数不更新时不能只盯着损失曲线猜应该直接去看梯度数值。total_norm 0.0 for p in model.parameters(): if p.grad is not None: param_norm p.grad.data.norm(2) total_norm param_norm.item() ** 2 total_norm total_norm ** 0.5 print(f全局梯度范数: {total_norm:.4f})如果梯度范数接近0大概率是梯度消失可以考虑换激活函数比如把sigmoid换成relu或gelu、调整网络初始化方式、使用残差连接。如果梯度范数指数级增长甚至变成NaN那多半是梯度爆炸最简单的手段是加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)梯度裁剪可以放在loss.backward()之后、optimizer.step()之前。它不会改变梯度方向只是把梯度范数限制在给定范围内相当于给更新踩了一脚刹车。还有一个经常被忽视的问题某些层的梯度一直为None。这说明反向传播根本没走到那层检查一下前向过程里是否不小心用了detach()或no_grad()或者输入数据是不是全部一样导致某个分支没有有效梯度。打印每一层参数名和对应p.grad是否为None很快就能定位。6. 自定义梯度与调试技巧hook、autograd.Function和数值校验6.1 register_hook查看和修改中间梯度反向传播过程中梯度流经每个张量节点。如果想看看某个中间层的梯度长什么样可以注册一个hookx torch.tensor(2.0, requires_gradTrue) y x ** 3 def print_grad(grad): print(中间梯度:, grad) return grad * 2 # 可以修改梯度 y.register_hook(print_grad) y.backward() print(x.grad) # 此时会多乘2register_hook的一个常见用途是梯度修改。比如想让某个中间层输出的梯度加速流动可以在这里手动乘以一个放大系数这也是很多论文里梯度缩放技巧的实现方式。另一个用途是调试在训练过程中打印梯度看哪一层的梯度异常。要注意的是hook优先于梯度累积执行也就是说x.grad最终累加的是hook返回之后的梯度。如果你在hook里直接返回了一个新张量它必须保持与传入梯度相同的形状否则后续计算会报错。6.2 自定义torch.autograd.Function有些算子PyTorch没有直接提供或者你想实现一个带特殊梯度规则的前向操作这时候可以继承torch.autograd.Function。核心是重写forward和backward两个静态方法class SquareAndScale(torch.autograd.Function): staticmethod def forward(ctx, input, scale): ctx.scale scale ctx.save_for_backward(input) return input * input * scale staticmethod def backward(ctx, grad_output): input, ctx.saved_tensors grad_input grad_output * 2 * input * ctx.scale grad_scale None return grad_input, grad_scale scale torch.tensor(2.0, requires_gradTrue) x torch.tensor(3.0, requires_gradTrue) y SquareAndScale.apply(x, scale) y.backward() print(x.grad) # 2 * 3 * 2 12这里forward里用save_for_backward保存了前向需要的数据backward里接收的grad_output是上游传来的梯度我们负责算出对输入的梯度。如果forward有多个输入backward必须返回对应数量的梯度某个输入不需要梯度时返回None。自定义Function常用于实现一些不可导或者梯度需要定制的前向过程比如带随机扰动的采样操作、量化操作等等。虽然现在很多场景都能用现有API组合解决但理解这个机制能让你对梯度计算的掌控力进入下一个层次。6.3 梯度检查用有限差分验证反向实现当你写完一个自定义算子很担心反向传播算得对不对的时候可以用数值梯度做校验。数值梯度的原理是导数的定义式$$\frac{\partial f}{\partial x} \approx \frac{f(x \epsilon) - f(x - \epsilon)}{2\epsilon}$$然后把它和autograd求出的梯度比较。PyTorch里已经封装好了现成的工具from torch.autograd import gradcheck # 定义一个函数接收一个输入 my_func SquareAndScale.apply x torch.randn(3, dtypetorch.double, requires_gradTrue) scale torch.randn(1, dtypetorch.double, requires_gradTrue) inputs (x, scale) test gradcheck(my_func, inputs, eps1e-6, atol1e-4) print(梯度校验通过:, test)这里必须把张量设置为torch.double类型因为数值微分对浮点精度很敏感单精度很容易因为舍入误差导致校验失败。gradcheck本身会采样多个随机点进行有限差分比较返回True说明你的反向传播实现和数值梯度吻合基本可以放心用。如果校验失败别急着改backward先检查forward是不是写错了再检查backward的返回顺序是否和forward的输入顺序一致这是最容易出错的地方。7. 写在最后我常用的梯度计算调试套路我在调模型时有一套固定排查流程。先看loss.backward()是否报错如果报错优先检查requires_grad、detach()、in-place操作这三项。不报错但模型不收敛时我会打印各层梯度范数确认梯度是否真的到达了所有需要更新的参数。梯度到了但数值混乱再考虑梯度裁剪和初始化问题。还有一个省心的习惯在写关键模型之前先用一个最小样例验证梯度。比如定义一个只有一层的线性模型输入一个固定张量手算一次梯度再用PyTorch对比。这个过程虽然简单但能提前暴露很多隐性问题比训练了半天才发现梯度链路断了要划算得多。梯度计算在你刚接触时像一堵墙真正理解之后会发现它不过是一套按规则运转的机制。多动手写写小实验多看看中间节点的grad_fn远比死记API有意义。
