PyTorch autograd核心机制:自动微分与反向传播保姆级解析
先声明一下这个“PyTorch6”大概率是标题笔误或者某个教程里的版本代号咱们不看版本号直接聊autograd这个核心引擎本身。题目标题里已经点得很透了自动微分与反向传播的核心原理而且用了“保姆级”这个词说明目标读者是那些刚入门PyTorch、对张量操作还停留在“data requires_grad backward”三件套、遇到梯度问题就一脸懵的人。这篇文章就是写给这样的人看的当然如果你已经用PyTorch写过不少模型但始终没把autograd底层机制捋顺那这篇文章同样能帮你补上这块拼图。1. PyTorch的autograd为什么值得单独学一遍1.1 自动微分不是“黑魔法”而是一套按链式法则运行的记账系统很多人从TensorFlow转过来或者刚接触PyTorch时会把autograd当成一个“只要调backward就能拿到梯度”的黑盒子。其实autograd的本质非常简单它是一套运行在前向传播过程中的记账系统。在你执行张量运算的那一刻autograd会同时做两件事——把计算结果存下来并把这次运算的“配方”记录到一张图里。这张图在PyTorch里被称为动态计算图。它的节点是张量Tensor边是产生这个张量的运算函数。举个例子z x y那么x和y是指向叶子的节点add运算产生一个z节点同时z会持有一个grad_fn指向AddBackward对象。这个对象记录了它需要哪些输入来决定如何反向传播梯度。等你调用z.backward()时autograd就会沿着grad_fn链反向走一圈把梯度一路传回x和y。这套机制在概念上和大学高等数学里的链式法则一模一样如果z依赖yy依赖x那么dz/dx dz/dy * dy/dx。autograd做的所有事情就是递归地应用这个公式并且把每一步的偏导值都算出来乘在一起。它不聪明但它不会出错——只要你正确地表达了前向运算。1.2 为什么PyTorch选“动态图”而不是“静态图”你如果看过TensorFlow 1.x时代的老代码就会知道那时候要先定义一个静态计算图然后在一个会话里反复执行图结构一旦定义就不能随便改。PyTorch走的是动态图路线图是即用即建的模型前向跑一遍图就建出来一次下一次前向图又重新建一次。动态图最大的好处是调试直观。你可以在前向传播的任意位置打印中间张量、打断点甚至用Python的if和for来动态控制网络结构。这在处理变长序列、树结构数据、或者需要在运行时调整分支的模型时尤其方便。静态图的编译优化确实能带来性能收益但对于绝大多数研究者来说开发的便捷性远比那点性能提升更重要。这也是PyTorch能迅速在学术界站稳脚跟的核心原因之一——你的研究思路本身就是动态的PyTorch能跟上你的思路。1.3 autograd的未来趋势和现状顺便多说一句PyTorch团队这几年的方向已经不只是追求动态图的便利他们在2.x版本里开始做torch.compile把动态图在运行时编译成更高效的静态执行计划。但这不改变autograd的基础工作流——编译只是把“每次重新建图”变成“同一张图只建一次再用很多次”梯度计算依然依赖那一整套记账系统。所以你先把autograd搞懂后面学torch.compile、torch.func这些进阶功能时会顺畅得多。2. 核心机制拆解一张图看懂反向传播的记账系统2.1 requires_grad给数据贴上“我需要梯度”的标签想理解autograd第一个必须掌握的概念就是requires_grad。一个张量只要设置了requires_gradTruePyTorch就会把它当作计算图中的叶子节点并为它建立一套梯度记账结构。代码看一眼就懂了import torch x torch.tensor(2.0, requires_gradTrue) y x ** 2这里的x就是一个叶子张量它有requires_gradTrue意味着我想计算关于x的梯度。y是x**2的结果它也会有requires_gradTrue并且y的grad_fn指向PowBackward0。当你执行y.backward()时PyTorch会从y出发找到它是通过Pow运算生成的而Pow运算的输入是x于是它用Pow的导数公式计算出梯度dy/dx 2*x 4.0然后把这个值累积到x.grad里。这里有个很重要的细节只有当张量参与运算并生成了一个带grad_fn的结果时梯度才能传递。纯粹定义一个requires_gradTrue的张量但不拿它做任何运算那它就没有grad_fn相关的传播路径。新手最容易犯的错误之一是给所有输入都加了requires_gradTrue但忘了中间结果其实会自动带上记录或者反过来中间结果被重新赋值后梯度路径断了。2.2 grad_fn记录“我是怎么算出来的”grad_fn是autograd记账系统的核心。每个由运算生成的张量都有grad_fn它是反向传播时用于计算梯度的函数对象。前向传播网络中的张量依赖关系全部以grad_fn的形式串联起来。看这个例子import torch x torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) w torch.tensor([2.0, 3.0, 4.0], requires_gradTrue) a x * w # MulBackward0 b a.sum() # SumBackward0 print(a.grad_fn) # MulBackward0 object at ... print(b.grad_fn) # SumBackward0 object at ...b的grad_fn是SumBackward0它知道自己在反向传播时要怎么处理梯度——对一个求和操作来说梯度就是直接把上游梯度复制到每个被加元素上。a的grad_fn是MulBackward0它知道对乘法a x * w来说梯度应该分别乘以对方的系数。执行b.backward()后autograd会按b - a - x, w的路径反向遍历。这就是反向传播的全过程。你不用担心记不住每个运算对应的导数PyTorch内置了几乎所有常见运算的反向规则。真正需要你操心的是你自己设计了一套复杂的前向逻辑但没有按照autograd能理解的算子来组合它。2.3 叶子节点与非叶子节点两个需要区分清楚的概念继续往前走你还得搞清楚一个概念叶子节点和非叶子节点。简单来说你直接创建并且requires_gradTrue的张量是叶子节点由其他张量计算得到的结果是非叶子节点。上面例子里a就是非叶子节点x和w是叶子节点。两者最重要的区别在于反向传播之后叶子节点的.grad会保留梯度值非叶子节点的.grad默认不会保留。为什么内存优化。一个复杂的神经网络里有几百万个中间张量如果每个的梯度都保存下来显存早就爆了。所以PyTorch默认只保存叶子节点的梯度。但调试的时候你就想知道某个中间张量的梯度。这时候有几种办法在非叶子节点上调用.retain_grad()让PyTorch把它的梯度也保存下来。用register_hook为张量注册一个钩子函数在梯度计算出来的时候把它捕获。retain_grad示例x torch.tensor(2.0, requires_gradTrue) y torch.tensor(3.0, requires_gradTrue) z x * y z.retain_grad() # 手动保留中间结果 z 的梯度 loss z ** 2 loss.backward() print(z.grad) # 现在有值了这招在排查某个隐藏层梯度过大或过小时非常有用不用改代码结构一行就能看到中间梯度。2.4 链式法则在代码里的具体体现链式法则是反向传播的数学基础autograd就是它的工程实现。你不需要在写网络时手动推导每个函数相对于每个参数的偏导但这些偏导在backward时都会被自动执行。对神经网络来说假设有一个两层线性网络输入 x - 线性层1(w1, b1) - 激活函数 - 线性层2(w2, b2) - 损失 loss那么loss对w1的梯度需要通过loss对输出层的梯度、再乘以激活函数的导数、再乘以w2的系数逐步传回来。这个过程就是链式法则。每层反向传播的梯度都会作为“上游梯度”传入下一层。你可以把这个过程类比为一条生产流水线的逆向追溯最终产品的质量问题要追溯到每一步加工环节但追溯效率最高的方式是从成品反向往前找而不是从源头正向排查。反向传播之所以叫“反向”就是因为它从loss开始沿着前向运算的反方向一次性地算出所有参数的梯度。2.5 动态计算图的生命周期用一次就没了吗一个容易踩坑的点动态计算图在backward()之后会被释放。默认情况下你调用一次backward()autograd就会把这张图的内存清掉为下一次前向计算腾出空间。所以你如果在同一个计算图上想调两次backward()比如需要计算二阶梯度或者做某些对抗训练场景就得在第一次调用时设置retain_graphTrue。x torch.tensor(2.0, requires_gradTrue) y x ** 3 y.backward(retain_graphTrue) # 第一次反向传播保留图 print(x.grad) # 12.0 x.grad.zero_() # 清掉上一次的梯度 y.backward() # 第二次反向传播 print(x.grad) # 12.0注意即使第二次调用后没有设置retain_graph如果后面不再调用backward那也没事。但如果你在循环里反复对一个图backward而忘了设置retain_graphTrue会报RuntimeError。这个坑在后面常见问题章节还会细讲。3. 从零实操用autograd手写线性回归3.1 准备数据与初始化参数看了一堆原理总得上手跑一遍。这里我带你用autograd手写一个最简单的线性回归模型目标是把所有概念落在地上。这个例子足够小但autograd的每一个关键动作都会用到。先生成一组带噪声的线性数据import torch import matplotlib.pyplot as plt # 生成数据y 2x 3 噪声 torch.manual_seed(42) x_data torch.linspace(-1, 1, 100).reshape(-1, 1) y_true 2 * x_data 3 y_data y_true 0.2 * torch.randn_like(x_data)接下来初始化模型参数。注意这里是手动定义两个张量作为模型的权重和偏置而不是用nn.Linear因为我想让你看到梯度更新的完整闭环w torch.randn(1, requires_gradTrue) b torch.zeros(1, requires_gradTrue) print(w.grad) # 初始为 None因为还没进行反向传播requires_gradTrue表示w和b是我们要学习的参数autograd会跟踪它们的所有运算。3.2 前向传播、损失函数与backward手写前向传播def forward(x): return w * x b def mse_loss(pred, target): return ((pred - target) ** 2).mean()训练循环的核心步骤learning_rate 0.1 epochs 200 for epoch in range(epochs): # 前向传播 y_pred forward(x_data) loss mse_loss(y_pred, y_data) # 反向传播 loss.backward() # 手动更新参数禁止梯度跟踪 with torch.no_grad(): w - learning_rate * w.grad b - learning_rate * b.grad # 清空梯度防止累积 w.grad.zero_() b.grad.zero_() if (epoch 1) % 20 0: print(fEpoch {epoch1}/{epochs}, Loss: {loss.item():.4f}, w: {w.item():.4f}, b: {b.item():.4f})这段代码浓缩了autograd的全部精华。有几个点值得重点解释。第一loss.backward()执行时PyTorch会从loss出发按计算图的反向路径把梯度送到每一个requires_gradTrue的张量上。执行完后w.grad和b.grad就分别保存了损失函数对它们的偏导数。第二参数更新必须放在torch.no_grad()的上下文里。为什么因为如果在带梯度的环境下执行w w - learning_rate * w.grad这个减法运算又会生成新的计算图分支导致你下一次backward路径错乱或者计算图越滚越大。PyTorch设计哲学是即使你在做参数更新也不希望autograd追踪这些操作。所以用no_grad()把这步包起来。第三梯度必须清零。如果不调zero_()那么下一轮backward()时梯度会在已有值上累加。你会看到loss不降反升或者参数更新幅度越来越大最后直接发散。最终训练结束后你应该看到w收敛到2附近b收敛到3附近噪声影响下略有浮动。3.3 用optimizer替代手动更新实际写模型时你大概率会直接用torch.optim.SGD或者torch.optim.Adam来处理参数更新。替换方式很简单import torch.optim as optim w torch.randn(1, requires_gradTrue) b torch.zeros(1, requires_gradTrue) optimizer optim.SGD([w, b], lr0.1) for epoch in range(epochs): optimizer.zero_grad() y_pred forward(x_data) loss mse_loss(y_pred, y_data) loss.backward() optimizer.step()注意这里optimizer.zero_grad()必须放在backward()之前不能在之后。它的作用等同于手动调w.grad.zero_()和b.grad.zero_()。如果你用optimizer.step()之前忘记zero_grad()梯度会跨batch累加最终导致收敛失败。很多新手都会踩这个顺序问题的坑。3.4 手动实现和nn.Module的关系有了上面的基础你再去看nn.Linear和nn.Module的实现思路就很清晰了。nn.Linear内部就是定义了一组有requires_gradTrue的权重和偏置然后在前向传播时执行weight x bias。nn.Module的parameters()方法就是把这些张量统一管理起来方便你交给optimizer。所以当你定义一个模型时import torch.nn as nn class LinearRegression(nn.Module): def __init__(self): super().__init__() self.linear nn.Linear(1, 1) def forward(self, x): return self.linear(x)在训练里做的事情和我上面手写的逻辑完全一样只是封装层次更高了而已。理解这一点你就不会被nn.Module的神秘感唬住。3.5 验证梯度是否正确的小技巧在写完一个自定义结构后怎么确认梯度计算是对的有一个非常实用的小技巧——数值梯度检验gradient check。思路就是利用导数的定义f(x) ≈ (f(x eps) - f(x - eps)) / (2 * eps)然后和autograd算出来的梯度做对比。如果两者误差在1e-6量级内说明你的前向逻辑是正确的。PyTorch的torch.autograd.gradcheck接口可以直接做这件事但在简单场景下手动验证也很方便eps 1e-6 x_tensor torch.tensor(2.0, requires_gradTrue) y x_tensor ** 3 y.backward() autograd_grad x_tensor.grad.item() # 数值梯度 x1 2.0 eps x2 2.0 - eps numerical_grad ((x1 ** 3) - (x2 ** 3)) / (2 * eps) print(fautograd: {autograd_grad:.8f}, numerical: {numerical_grad:.8f})如果你写了一个自定义Function这个技巧能帮你快速定位前向或者反向实现里的bug。我的习惯是任何自定义算子先过一遍gradcheck再接入网络否则网络一深根本分不清是梯度算错了还是网络结构有问题。4. autograd常见报错和排查锦集4.1 梯度为None参数死活不更新这是新手遇到的最高频问题。查了一圈发现loss在变但是model.weight.grad一直是None参数更新毫无动静。原因通常是下面几种。第一张量没有设置requires_gradTrue。如果你手动创建参数张量但又忘了加这个参数那么autograd根本不会为它计算梯度。解决办法就是检查参数定义处确保requires_gradTrue。第二模型在torch.no_grad()上下文里执行了前向传播。我见过有人为了“加速”把整个训练循环包进了no_grad()里那当然不会有任何梯度。第三损失函数对应的路径上有一个detach()调用。detach()会切断计算图的连接就像在流水线上拆掉了一段传送带后面的梯度传不回来了。排查建议在loss.backward()之后用print(model.weight.grad)确认一下有没有梯度。如果确实是None再往前检查数据、模型、损失函数看哪一步把计算图切断了。4.2 RuntimeError: a leaf Variable that requires grad has been used in an in-place operation这个报错也极其经典。出现的原因是在一个requires_gradTrue的叶子张量上执行了原地操作比如x 1、x.zero_()、weight.data.add_(lr * grad)之类的。PyTorch发出这个错误是因为原地修改有可能会覆盖掉autograd反向传播时需要引用的原始值导致梯度计算错误。解决方案分两种如果你只是想清空梯度使用x.grad.zero_()这个操作针对的是.grad不会触碰数据本身所以合法。如果是想更新参数稳妥的做法是x.data - lr * x.grad或者with torch.no_grad(): x - ...。如果你确实需要对中间张量做原地操作可以用x.clone()先复制一份再改。但一般来说训练循环里的参数更新交给optimizer处理不会碰到这个错误。4.3 想对同一个损失调用两次backward结果报错前文提到动态计算图在第一次backward()之后就被释放了。如果你需要对同一个计算图调用两次backward()比如算二阶梯度、或者某些多任务梯度操作必须在第一次调用时加上retain_graphTrueloss.backward(retain_graphTrue)更推荐的做法是如果需要多次反向传播尽量重新构造前向过程。重算一遍前向的成本通常比一直占着计算图维护内存要划算。我在做对抗训练Adversarial Training时经常需要先对对抗样本算一次梯度再对干净样本算一次梯度每次都重新前向一次既清晰又不容易出错。4.4 对非标量loss调用backward时没有传gradient参数默认情况下backward()要求损失是一个标量。如果你对向量或者矩阵直接调用backward()比如x torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) y x ** 2 y.backward() # 报错PyTorch会告诉你只能对标量输出进行backward。为什么因为对向量求梯度你需要先明确“这个向量的哪个分量需要梯度信息”。数学上只有标量函数对变量才有唯一的梯度向量对变量的梯度是一个雅可比矩阵不能直接以向量的形式返回。如果你确实需要对非标量调用必须传入一个与y形状一致的gradient参数作为权重y.backward(torch.tensor([1.0, 1.0, 1.0]))这行代码等价于先对y的所有元素求和再对这个和求梯度。在实践中最常遇到这个场景的是模型输出batch_size个预测结果你用loss torch.nn.functional.mse_loss(pred, target, reductionnone)得到了一个向量损失然后试图直接对这个向量调用backward()。正确做法是先对向量求均值或求和让loss变成标量再backward。4.5 跨设备/跨精度的梯度路径断裂还有一个容易被忽略的问题做了.to(device)之后可能导致梯度路径失效。比如你创建了一个requires_gradTrue的cpu张量然后通过.cpu()或者.to(cpu)把它挪到CPU如果它已经在CPU上那没问题但如果从CUDA挪到CPU梯度不会跟着传回去。PyTorch的autograd只支持同一设备上的计算图跨设备操作会直接切断梯度的传播。相似地跨精度也可能有问题。如果你把float32的参数和float16的数据混着做运算某些运算在反向传播时可能不稳定甚至直接报错。我建议模型参数和输入数据的dtype保持一致必要的时候用.to(device).to(dtype)做明确转换。5. autograd的高阶用法与性能优化5.1 用torch.no_grad()节省显存与加速验证集评估和推理阶段我们不需要计算梯度所以应该把整个推断过程包在torch.no_grad()上下文里。这样做的好处有两个一是省去autograd记账的开销推理速度更快二是不用保存中间激活值显存占用大幅下降。model.eval() with torch.no_grad(): val_pred model(val_data)你可能会觉得“反正模型不大不差这点显存”但当你用大模型或者处理长序列时省下的显存可能够你把batch size翻一倍。我在做Transformer相关的任务时验证集推断如果不包no_grad()一个batch的显存占用能差出2-3倍。养成这个习惯真的很重要。5.2 梯度累积用小显存模拟大batch如果你需要用256的batch size但显存只够跑64梯度累积Gradient Accumulation是一个经典解法。思路是先用小batch计算梯度但不清空梯度连续累积若干次后再更新参数。accumulation_steps 4 optimizer.zero_grad() for i, data in enumerate(train_loader): loss compute_loss(data) loss loss / accumulation_steps # 归一化确保总梯度相当于大batch做了平均 loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这里有个关键点需要把每个小batch的loss除以accumulation_steps否则梯度累积下来会把有效学习率放大好几倍导致训练不稳定。这个除法本质上是用数学上的链式法则均匀分摊了梯度信息autograd不关心你除以多少它只会忠实地把梯度传回给每个参数。5.3 自定义autograd.Function打破内置算子的边界PyTorch官方提供的算子覆盖了绝大多数场景但科研中总会碰到实现自定义操作的时刻。这时候你要继承torch.autograd.Function手动实现forward和backwardclass MyCustomOp(torch.autograd.Function): staticmethod def forward(ctx, input): ctx.save_for_backward(input) # 保存反向所需的数据 return input * input staticmethod def backward(ctx, grad_output): input, ctx.saved_tensors return grad_output * 2 * input x torch.tensor(3.0, requires_gradTrue) y MyCustomOp.apply(x) y.backward() print(x.grad) # 6.0写自定义Function的时候最容易出bug的地方是backward的输入输出数量必须和forward的输入输出数量严格对应。grad_output是上游传来的梯度你要在backward里把它乘上你这一步的局部导数。ctx.save_for_backward是用来保存前向时的中间量因为你可能在反向时还要用到它们。我刚才提到的torch.autograd.gradcheck就是给这种自定义算子准备的。5.4 梯度裁剪、梯度累积的位置和先后顺序做RNN或者Transformer训练时经常需要梯度裁剪防止梯度爆炸。记住顺序是backward()-clip_grad_norm_()-optimizer.step()。有些新手会在optimizer.zero_grad()之后、backward()之前调用裁剪那当然没用因为那时候梯度还是None。loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step()这里的clip_grad_norm_是对所有参数的梯度求一个全局范数然后按比例缩放防止某个参数梯度过大导致更新步长异常。这也是autograd计算完梯度之后、优化器更新参数之前的常规操作。5.5 关于PyTorch版本和环境的一点提醒前面说过标题里的“PyTorch6”大概率是笔误但我还是想提醒一句不同PyTorch版本在autograd的实现细节上偶尔有小改动比如torch.autograd.graph里新增的Node接口、torch.autograd.Function的vmap支持等。你学习的时候一定要先查清楚自己环境里的版本。看官方API文档时左上角有版本选择下拉框根据实际安装版本查看对应文档能省掉很多因为API变化导致的迷惑。顺带说一句环境问题很多人在Windows上用Anaconda装PyTorch遇到OSError: [WinError 1114]之类的DLL加载失败这通常不是代码问题而是CUDA版本和PyTorch版本不匹配、或者Visual C运行库缺失导致的。如果你暂时只想跑通autograd代码用CPU版本的PyTorch完全可以autograd机制本身不依赖GPU。等后面需要跑大模型了再按官方教程重新安装匹配的CUDA版本。6. 几个比较少人提但很有价值的小技巧6.1 用.f.grad和.grad.data的区别优化器更新参数时常见两种写法。一种是w.grad.data.zero_()另一种是w.grad.zero_()。从PyTorch行为上看效果几乎一样但源码里optimizer大量使用了param.grad的原地操作。tensor.data返回的是一个和原张量共享存储空间但没有requires_grad属性的新张量操作时不会被autograd追踪所以用来更新数据很安全。如果你在自定义训练循环里想对梯度做特殊处理比如只更新梯度中的正数部分推荐先clone一下梯度再操作grad_copy w.grad.clone()。直接修改w.grad会影响后续所有依赖它的逻辑包括梯度的二次计算容易出隐蔽的bug。6.2 启用异常检测快速定位梯度路径问题PyTorch提供了一个运行时开关torch.autograd.set_detect_anomaly(True)。打开之后autograd会在前向传播时额外检查计算图一旦检测到某个操作可能导致反向传播失败它会尽量在报错时指出具体是哪个前向操作引发的。这个开关会显著拖慢训练速度所以我只在调试阶段用torch.autograd.set_detect_anomaly(True)调试完马上关掉。它最大的价值在于当你在深层网络里收到“one of the variables needed for gradient computation has been modified by an inplace operation”这种报错时它能直接告诉你问题出在哪一行前向代码省去你一层一层猜。6.3 断点调试时怎么看梯度在PyCharm或者VS Code里调试训练循环时你可以在loss.backward()之后打断点观察model.weight.grad的形状和值。如果梯度全部为nan大概率是前向运算里出现了除零或者数值溢出如果梯度为全零可能是激活函数在饱和区比如sigmoid输入绝对值太大如果梯度爆炸去看是不是没做梯度裁剪。实际调试网络时除了看参数梯度偶尔也要看激活值分布和loss曲线。梯度只是信号的“流量”激活值的范围才决定了流量是否稳定。这些信息结合起来才能真正定位问题。6.4 理解.autograd不需要背每个算子的导数公式不少学员问我autograd底层是不是已经把每个算子的导数公式都内置了是的而且你不需要背下它们。你只需要掌握两个能力第一能判断自己的前向代码里哪些张量是“需要梯度”的第二能沿计算图推理出一条梯度传播路径在梯度不更新或更新异常时迅速定位切断点。这两点这篇文章里都讲透了。7. 写在最后我的一点使用体感从PyTorch 0.4开始用到现在autograd给我的最大感受是它把“求导”这件事中最繁琐、最易犯错的部分拿走了但你依然得知道梯度从哪来、到哪去。就像开车不用懂内燃机每个零件的原理但你必须知道什么时候踩油门、什么时候踩刹车、仪表盘报警灯亮起来时往哪看。autograd对你来说就是那台发动机你不需要亲手造它但要想把车开快还是得了解它的脾气。最后分享一个小技巧当你要写一个复杂模型而又不确定梯度路线是否正确时先把模型输入设成全requires_gradTrue的小向量跑一次backward()然后遍历model.parameters()打印每层梯度的范数。如果整个链路里有任何一层梯度是None你立刻就能发现中断点在哪这比等训练好几个epoch之后才发现参数没更新要高效得多。这套排查流程我每次写新结构时都会跑一遍几乎成了肌肉记忆。希望对你有用。