1. 从一次显存泄漏说起为什么必须搞懂计算图刚接触PyTorch那会儿我写过一个看起来毫无问题的训练循环每个batch前向传播、算loss、反向传播、更新参数然后进入下一轮。代码跑起来loss确实在降但显存占用却像爬楼梯一样一路往上走跑了几百个step之后直接OOM。当时我第一反应是数据没释放干净于是各种del、gc.collect()、torch.cuda.empty_cache()全招呼上结果一点用没有。后来把模型里一段调试用的中间变量打印去掉显存立刻稳住了。问题就出在那行打印上——它把一个带梯度的中间张量保存了下来而这个张量背后挂着整张计算图图不释放显存自然不回收。这件事让我意识到PyTorch里很多玄学问题根子都在计算图Computational Graph上。你不需要成为框架开发者但只要你在用PyTorch训练模型计算图就决定了你的显存怎么分配、梯度怎么流动、哪些操作能求导哪些不能、什么时候该用detach()、为什么有些in-place操作会报错。这些不是边角知识而是每天都会撞上的东西。这篇内容我想把计算图这件事从头讲透。不是那种计算图就是节点和边的教科书定义而是从实际写代码的角度出发动态图到底动态在哪、反向传播时图是怎么被消费的、叶子节点和非叶子节点的区别为什么会导致梯度拿不到、retain_graph什么时候必须开、torch.no_grad()和detach()到底差在哪。适合已经能跑通训练脚本、但遇到报错或显存问题只能靠搜索和试错的读者。看完之后你至少能做到看到报错信息就知道是哪类图的问题而不是盲目加torch.cuda.empty_cache()。2. 动态图的动态到底体现在哪里2.1 前向传播时图才被现场搭建PyTorch的计算图是动态的这句话很多人会背但真正理解它意味着什么的人不多。我用一个对比来说明。静态图框架比如早期的一些方案的思路是你先定义好整个计算流程框架把它编译成一张固定的图然后你往里面喂数据。图的结构在运行前就定死了循环、条件分支都得用框架提供的特殊语法来写。PyTorch完全不是这个路子。你写的是普通的Python代码for循环、if判断、甚至递归调用都随便用。每执行一行涉及张量的运算PyTorch就在背后记录一个节点。前向传播跑完图也就搭完了。下一次迭代代码重新执行一遍图重新搭一遍。这就是动态的含义——图的生命周期和一次前向传播绑定。这个特性带来的直接好处是调试极其方便。你可以在任意一行print张量的形状和值可以用Python的调试器单步跟踪可以随时改变网络结构。代价是每次迭代都有建图的开销不过对于绝大多数场景这点开销相比GPU计算可以忽略。我举个实际会遇到的例子。假设你要实现一个带early stopping的循环网络根据某个条件决定是否继续展开时间步def forward(self, x, max_steps10): h self.init_hidden(x) outputs [] for step in range(max_steps): h self.cell(x, h) outputs.append(h) if h.abs().mean() 1e-3: # 动态条件 break return torch.stack(outputs)这段代码在动态图框架里直接就能跑break什么时候触发取决于运行时的数据。如果换成静态图你得用框架提供的while_loop之类的原语来写麻烦得多。这就是动态图最直观的价值。2.2 每个张量身上的三个关键属性要理解计算图得先搞清楚一个张量身上挂了哪些和图相关的信息。我把它归纳成三个属性理解了这三个后面很多问题都能自己推导出来。第一个是requires_grad。这个布尔值决定这个张量是否需要计算梯度。默认情况下你手动创建的张量requires_gradFalse而模型参数nn.Parameter默认是True。只要一个运算的任意输入requires_gradTrue输出就会requires_gradTrue并且会在图中记录这个运算。第二个是grad_fn。这是指向创建该张量的函数的引用。如果你看到一个张量的grad_fn是None说明它是叶子节点不是由某个运算产生的。如果grad_fn有值比如AddBackward0说明它是由加法产生的反向传播时就会调用对应的反向函数。第三个是is_leaf。叶子节点是用户直接创建的、或者requires_gradFalse的张量。模型参数是典型的叶子节点。非叶子节点是由运算产生的中间结果。这个区分极其重要因为只有叶子节点的梯度会被默认保留在.grad属性里非叶子节点的梯度在反向传播后默认不保存。import torch x torch.tensor([1.0, 2.0], requires_gradTrue) # 叶子节点 y x * 2 # 非叶子节点 z y.sum() # 非叶子节点 print(x.is_leaf) # True print(y.is_leaf) # False print(y.grad_fn) # MulBackward0 object z.backward() print(x.grad) # tensor([2., 2.]) 叶子节点梯度保留 print(y.grad) # None 非叶子节点梯度默认不保留很多人第一次遇到为什么我拿不到中间层的梯度就是栽在这里。想拿非叶子节点的梯度得用y.retain_grad()显式告诉PyTorch保留它。2.3 图是用完即焚的动态图还有一个容易被忽略的特性反向传播完成后图默认就被释放了。这是为了省内存。你想想一个深层网络前向传播会产生大量中间激活值这些值在反向传播时都要用到所以必须存着。一旦反向传播算完这些中间值就没用了图也就没存在的必要了。这个设计导致一个常见报错你想对同一个loss做两次backward()第二次就会报Trying to backward through the graph a second time。因为第一次反向传播已经把图消费掉了。解决办法是第一次调用时加retain_graphTrue让图保留下来。但要注意保留图意味着中间激活值不释放显存占用会显著增加所以除非确实需要比如多任务loss分别反向否则不要随便开。我在实际项目里见过有人为了保险每次都加retain_graphTrue结果显存直接翻倍。这个参数不是用来防报错的它有明确的使用场景。3. 反向传播时图是怎么被消费的3.1 从loss到参数的梯度链路反向传播的本质是链式法则。计算图记录了前向传播的每一步运算反向传播就是从loss这个标量出发沿着图反向走一遍每经过一个节点就计算该节点的局部梯度然后乘上从上游传下来的梯度得到传给下游的梯度。用个简单的例子手动推一遍。假设x torch.tensor(3.0, requires_gradTrue) w torch.tensor(2.0, requires_gradTrue) b torch.tensor(1.0, requires_gradTrue) y w * x b # y 2*3 1 7 loss y ** 2 # loss 49前向传播建的图是x和w经过乘法得到wxwx和b经过加法得到yy经过平方得到loss。反向传播时loss y^2所以d(loss)/d(y) 2y 14y wx b所以d(y)/d(wx) 1d(y)/d(b) 1wx w * x所以d(wx)/d(w) x 3d(wx)/d(x) w 2链式相乘d(loss)/d(w) 14 * 1 * 3 42d(loss)/d(b) 14 * 1 14d(loss)/d(x) 14 * 1 * 2 28。PyTorch的grad_fn就是干这个的。每个grad_fn对象知道两件事怎么计算自己的局部梯度以及怎么把上游梯度传下去。backward()调用时从loss的grad_fn开始按拓扑逆序依次调用每个节点的反向函数。3.2 为什么梯度会累加而不是覆盖这是新手最容易踩的坑之一。PyTorch的梯度是累加的不是覆盖的。也就是说如果你连续做两次反向传播而没有清零梯度第二次的梯度会加到第一次上面。x torch.tensor([1.0], requires_gradTrue) y1 x * 2 y1.backward() print(x.grad) # tensor([2.]) y2 x * 3 y2.backward() print(x.grad) # tensor([5.]) 2 3不是3这个设计是有意为之的。在某些场景下比如RNN的梯度累积、多loss加权求和累加是需要的。但对于标准的训练循环你必须在每次反向传播前调用optimizer.zero_grad()否则梯度会越滚越大训练直接发散。我见过有人写训练循环时把zero_grad()放在backward()后面逻辑上好像也对——这次用完清掉下次用的时候是干净的。但这样写有个隐患如果某次迭代因为异常跳过了backward()梯度就没被清掉会污染下一次。标准写法是放在backward()之前每次迭代开头就清。3.3 梯度累积的实战用法既然梯度是累加的那就可以利用这个特性做梯度累积。当显存不够、没法开大batch时可以把一个大batch拆成几个小batch分别前向反向梯度自然累加最后再统一更新参数。效果等价于大batch训练。accumulation_steps 4 optimizer.zero_grad() for i, (inputs, labels) in enumerate(dataloader): outputs model(inputs) loss criterion(outputs, labels) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意loss要除以accumulation_steps否则累加后的梯度是原来的N倍相当于学习率放大了N倍。这个细节很多人会漏掉导致训练不稳定。提示梯度累积期间计算图在每个小batch反向传播后就被释放了所以显存占用和单个小batch相当不会因为累积而增加。这是它相比直接开大batch的优势。4. 那些和图相关的报错根因都在哪4.1 element 0 of tensors does not require grad这个报错的意思是你调用了backward()但loss这个张量的requires_grad是False没有图可以反向传播。常见原因有几个。第一种模型参数被冻结了。你可能手动把某些层的requires_grad设成了False或者用了torch.no_grad()包裹了前向传播。检查一下loss的requires_grad属性如果是False往上追溯是哪个环节断开了。第二种loss计算时用了不可导的操作。比如你把loss转成了Python float或者用了.item()或者经过了一个requires_gradFalse的张量运算。典型的是用accuracy之类的指标当lossaccuracy是离散的不可导。第三种输入数据没有requires_grad而且模型参数也被冻结。这种情况图根本建不起来。排查思路很简单从loss开始打印它的grad_fn和requires_grad然后往前找第一个grad_fn为None的地方那里就是断点。4.2 Trying to backward through the graph a second time前面提过图默认用完即焚。这个报错就是你试图对同一张图做第二次反向传播。除了加retain_graphTrue更值得思考的是你真的需要两次反向吗需要两次反向的典型场景是GAN。生成器和判别器的loss是分开的但都依赖于判别器对生成样本的输出。如果先算判别器loss反向图就被消费了再算生成器loss就没图了。这时候要么用retain_graphTrue要么把两个loss合在一起反向。# 方案一合并loss d_loss ... g_loss ... total_loss d_loss g_loss total_loss.backward() # 方案二retain_graph d_loss.backward(retain_graphTrue) g_loss.backward()方案一更省显存但两个loss的梯度会混在一起如果需要对不同参数做不同处理就不方便。方案二更灵活但显存占用高。实际用哪个取决于你的优化策略。还有一种情况是你在循环里反复对同一个loss反向这通常是代码逻辑写错了比如把backward()写在了内层循环里。这种要改代码结构而不是加retain_graph硬扛。4.3 in-place操作导致的版本冲突one of the variables needed for gradient computation has been modified by an inplace operation——这个报错也很常见。原因是反向传播需要用到前向传播时的某个中间值但你在前向之后、反向之前用in-place操作把这个值改了PyTorch检测到版本号不一致就报错了。x torch.tensor([1.0, 2.0], requires_gradTrue) y x * 2 y 1 # in-place操作y的版本号变了 z y.sum() z.backward() # 报错PyTorch用版本号机制来检测这种情况。每个张量有个_version属性in-place操作会递增它。反向传播时如果发现版本号对不上就知道中间值被篡改了。解决办法是用out-of-place操作替代y y 1 # 创建新张量不修改原张量或者用y y.clone() 1先克隆再操作。但克隆会增加显存所以优先用out-of-place写法。有些in-place操作是安全的比如relu_()在某些情况下不会触发报错因为它的反向不需要前向的输入值。但为了代码可维护性我建议除非有明确的性能需求否则统一用out-of-place写法。省那点显存不值得调试成本。4.4 显存泄漏的排查链路回到开头那个显存泄漏的例子。排查这类问题的完整链路是这样的第一步确认是不是计算图没释放。在训练循环里打印torch.cuda.memory_allocated()看它是否随step单调递增。如果是基本可以确定有东西没释放。第二步检查有没有把带梯度的张量存进了列表或字典。比如你把每个step的loss存进一个listloss是带图的图就一直挂着。正确做法是存loss.item()只存数值。第三步检查有没有在日志或调试代码里保存了中间张量。比如writer.add_scalar(grad, model.layer.weight.grad, step)如果weight.grad没被detach它可能带着图。第四步检查有没有忘记zero_grad()。虽然这不会导致图泄漏但会让梯度越来越大间接增加显存。第五步如果以上都排除了用torch.autograd.set_detect_anomaly(True)开启异常检测它会告诉你哪个操作导致了问题。这个模式很慢只在调试时开。我自己的经验是90%的显存泄漏都是把带梯度的张量存起来了这一类。养成习惯任何要长期保存的张量先.detach()或者.item()。5. no_grad、detach和inference_mode的取舍5.1 三者的语义差异这三个东西经常被混用但语义完全不同。torch.no_grad()是一个上下文管理器进入这个上下文后所有新创建的张量requires_grad都是False不会建图。它影响的是上下文内部的所有运算。detach()是张量的方法它返回一个新张量和原张量共享数据但requires_gradFalse且从图中脱离。原张量不受影响。torch.inference_mode()是后来引入的比no_grad更激进。它不仅不建图还会关闭一些autograd的元数据跟踪性能更好。但它有个限制在inference_mode下创建的张量之后不能在需要梯度的计算中使用。x torch.tensor([1.0], requires_gradTrue) with torch.no_grad(): y x * 2 print(y.requires_grad) # False z x * 2 z_detached z.detach() print(z.requires_grad) # True print(z_detached.requires_grad) # False5.2 验证和推理阶段该用哪个模型验证和推理时你不需要梯度用no_grad或inference_mode都能省显存、加速。选哪个如果你的代码只在推理阶段用且不涉及后续需要梯度的操作用inference_mode它更快。如果推理结果可能被用于后续需要梯度的计算比如某些元学习场景用no_grad更安全。我个人的习惯是纯推理用inference_mode验证阶段用no_grad。因为验证阶段有时会算一些指标虽然不需要梯度但代码可能和训练阶段共享用no_grad兼容性更好。注意no_grad和inference_mode都不会影响已经建好的图。如果你在上下文外面创建了带梯度的张量在上下文里面用它做运算运算结果不带梯度但原张量的图还在。要彻底断开得用detach()。5.3 detach的典型使用场景detach()最常用的场景是你想用某个网络的输出作为另一个网络的输入但不希望梯度回传到第一个网络。比如在强化学习里你用target network计算目标值这个目标值不应该对target network求梯度只对online network求梯度。这时候就要对target network的输出做detach()。with torch.no_grad(): target_q target_network(next_state) # 或者 target_q target_network(next_state).detach()两种写法效果类似但detach()更灵活因为它可以在图已经建好之后再断开。比如你先用target network算了一堆东西最后才决定哪部分要detach。还有一个场景是截断时间序列的反向传播。RNN展开太长时显存扛不住可以在某个时间步对hidden state做detach()让梯度只回传有限步。for step in range(seq_len): output, hidden rnn(input[step], hidden) if step % truncate_every 0: hidden hidden.detach()这样梯度最多回传truncate_every步显存可控。6. 自定义Function与图的手动干预6.1 什么时候需要自定义反向绝大多数情况下PyTorch内置的运算已经覆盖了你的需求。但有些场景需要自定义反向你实现了一个新的算子或者你想在前向和反向之间做一些特殊处理或者你想用更高效的数值稳定实现。自定义Function需要继承torch.autograd.Function实现forward和backward两个静态方法。forward里做计算backward里根据上游梯度计算对输入的梯度。class MyReLU(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x.clamp(min0) staticmethod def backward(ctx, grad_output): x, ctx.saved_tensors grad_input grad_output.clone() grad_input[x 0] 0 return grad_input关键点是ctx.save_for_backward它把反向需要的前向中间值存起来。注意只能存张量存其他类型用ctx.other value。存张量必须用save_for_backward否则可能出问题。6.2 自定义Function的常见坑第一个坑是backward的返回值数量必须和forward的输入数量一致。如果某个输入不需要梯度返回None。第二个坑是backward里不要做in-place操作修改grad_output因为它可能被其他地方引用。要改先clone()。第三个坑是save_for_backward存的张量如果在前向之后被修改了反向时会报版本错误。所以存进去之后就别动它。第四个坑是forward里如果用了no_grad上下文save_for_backward仍然有效但要注意上下文的范围。我实际用自定义Function最多的场景是实现一些数值不稳定的loss。比如某些涉及log-sum-exp的操作直接用内置函数可能溢出自己写一个数值稳定的版本同时自定义反向保证梯度正确。6.3 用hook观察和干预梯度除了自定义FunctionPyTorch还提供了hook机制可以在不修改模型代码的情况下观察或修改梯度。register_hook可以注册在张量上反向传播经过这个张量时触发def print_grad(grad): print(grad.abs().mean()) x torch.tensor([1.0, 2.0], requires_gradTrue) y x * 2 y.register_hook(print_grad) z y.sum() z.backward()register_full_backward_hook可以注册在模块上观察模块输入输出的梯度。这在调试梯度消失或爆炸时很有用。for name, module in model.named_modules(): module.register_full_backward_hook( lambda m, gin, gout: print(f{name}: {gout[0].abs().mean()}) )hook的返回值如果非None会替换原来的梯度。这个特性可以用来做梯度裁剪、梯度缩放等操作。但要注意hook里修改梯度是in-place的可能影响其他hook。7. 计算图视角下的显存优化实践7.1 激活值重计算反向传播需要前向的中间激活值这些值占显存的大头。激活值重计算也叫gradient checkpointing的思路是前向传播时只存少数几个关键节点反向传播需要中间值时重新算一遍。用计算换显存。PyTorch提供了torch.utils.checkpointfrom torch.utils.checkpoint import checkpoint def forward(self, x): x checkpoint(self.block1, x) x checkpoint(self.block2, x) return x每个checkpoint包裹的模块前向时不存中间激活反向时重新前向一遍。显存能省很多代价是训练速度慢一些大概多一次前向的时间。用checkpoint有几个注意点。第一被包裹的模块输入必须requires_gradTrue否则checkpoint不生效。第二如果模块里有随机性比如dropout前向重算时结果可能不一致需要用preserve_rng_stateTrue。第三checkpoint对batch norm不太友好因为重算时统计量可能变了。7.2 混合精度与图的关系混合精度训练AMP用float16做前向和反向float32做参数更新。它和计算图的关系在于autograd在float16下也能正常工作但梯度可能下溢。AMP的GradScaler就是用来解决这个问题的——它把loss放大反向得到的梯度也放大更新前再缩回去。scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意scaler.scale(loss)返回的是一个新的loss张量它带着缩放系数。反向传播时梯度也被缩放scaler.step会先检查有没有inf/nan没有才更新参数。AMP下计算图的结构和float32一样只是数据类型变了。但要注意有些操作在float16下不稳定比如softmax、layer normautocast会自动把这些操作转成float32。你不需要手动干预但要知道这个机制存在。7.3 梯度检查点与no_grad的配合有个容易混淆的点checkpoint内部的前向重算是在no_grad下做的吗不是。重算时需要建图否则反向没法继续。checkpoint的实现是第一次前向在no_grad下做不存图反向时重新前向这次建图然后立即反向用完释放。所以checkpoint包裹的模块其内部不能有依赖全局状态的运算因为重算时状态可能变了。这也是为什么checkpoint对dropout、batchnorm要特别处理。我实际用checkpoint的经验是在Transformer类模型上效果最明显因为每层都有大量激活值。一般从中间几层开始加checkpoint前面几层激活值小加了收益不大后面几层如果接近输出也可以不加。具体加在哪用显存分析工具看哪层占用大就加哪层。8. 一些容易被忽略的边界情况8.1 叶子节点的梯度什么时候是None叶子节点默认保留梯度但有些情况下grad是None。第一种这个叶子节点没有参与loss计算。比如你定义了两个参数但loss只用了其中一个另一个的grad就是None。第二种这个叶子节点的requires_gradFalse。第三种反向传播还没执行。optimizer.step()会跳过grad为None的参数所以不会报错。但如果你手动更新参数要检查grad是否为None。8.2 高阶导数的图PyTorch支持高阶导数但默认只建一阶图。要算二阶导需要在反向传播时也建图x torch.tensor([2.0], requires_gradTrue) y x ** 3 dy torch.autograd.grad(y, x, create_graphTrue)[0] d2y torch.autograd.grad(dy, x)[0] print(d2y) # 6*x 12create_graphTrue让反向传播的过程也被记录从而可以再反向。高阶导数在元学习、物理模拟等场景有用但显存开销大因为图会嵌套。8.3 多GPU下的图DataParallel和DistributedDataParallel下每个GPU上的前向传播建各自的图反向传播时梯度汇总到主GPU。图本身不跨GPU每个GPU独立管理自己的图。所以显存优化要针对单卡来做多卡只是把batch拆开了。DDP比DP效率高的原因之一就是DDP在每个进程里独立建图通信只在梯度汇总时发生而DP有额外的图复制开销。8.4 图的可视化想直观看看图长什么样可以用torchviz或tensorboard的graph功能。torchviz最简单from torchviz import make_dot y model(x) make_dot(y, paramsdict(model.named_parameters())).render(graph, formatpng)生成的图会显示每个运算节点和梯度函数。对于复杂模型图会很大但看看关键路径还是有帮助的。我一般只在调试奇怪的反向问题时才用平时靠grad_fn和报错信息就够了。9. 我踩过的几个真实坑第一个坑是关于loss.item()的。早期我写日志时直接存loss没加.item()结果每个step的loss都带着图跑了几百步显存就爆了。后来改成loss.item()问题消失。这个坑的本质是Python的list会持有对象的引用带图的张量被list引用着图就不会释放。第二个坑是关于retain_graph的。有次做多任务学习两个loss分别反向我加了retain_graphTrue结果显存翻倍。后来改成把两个loss加起来一次反向显存恢复正常。多任务loss如果不需要分别处理梯度合并反向是最优解。第三个坑是关于in-place的。我在一个自定义模块里用了x residual训练时报版本错误。改成x x residual就好了。但这样会多一份显存对于大模型有影响。后来我改成在残差连接处用x x residual其他地方尽量用in-place平衡了显存和正确性。第四个坑是关于detach()的。有次做target network更新我忘了对target的输出做detach()结果梯度回传到了target networktarget network被意外更新了训练完全跑偏。这个bug很难发现因为loss看起来在降但效果就是不对。后来我养成了习惯任何不应该被更新的网络输出一律detach()。第五个坑是关于no_grad的。有次在验证阶段用了no_grad但验证代码里有一行把验证loss存进了全局变量训练阶段又用这个loss做了一些计算。因为验证loss是在no_grad下算的没有图训练阶段用的时候报错。这个问题的根源是训练和验证的代码耦合太紧后来我把它们彻底分开各自管理自己的计算图。这些坑的共同点是它们都不是PyTorch的bug而是对计算图生命周期理解不到位导致的。理解了图什么时候建、什么时候释放、什么时候断开这些问题都能自己推导出答案。计算图这个东西刚开始学的时候觉得抽象但用久了会发现它其实是PyTorch里最讲道理的部分。每一个报错都有明确的物理含义每一个参数都有清晰的使用场景。你不需要背下来所有细节但建立起图的生命周期这个心智模型之后遇到问题就知道往哪个方向查。我现在的习惯是看到任何autograd相关的报错先问自己三个问题图建起来了吗图被消费了吗图被意外保留了吗这三个问题能解决大部分问题。 
长河君 · 2026/10/15 18:10:33 · 阅读 1.2k PyTorch计算图深度解析：从显存泄漏到梯度优化的实战指南
1. 从一次显存泄漏说起为什么必须搞懂计算图刚接触PyTorch那会儿我写过一个看起来毫无问题的训练循环每个batch前向传播、算loss、反向传播、更新参数然后进入下一轮。代码跑起来loss确实在降但显存占用却像爬楼梯一样一路往上走跑了几百个step之后直接OOM。当时我第一反应是数据没释放干净于是各种del、gc.collect()、torch.cuda.empty_cache()全招呼上结果一点用没有。后来把模型里一段调试用的中间变量打印去掉显存立刻稳住了。问题就出在那行打印上——它把一个带梯度的中间张量保存了下来而这个张量背后挂着整张计算图图不释放显存自然不回收。这件事让我意识到PyTorch里很多玄学问题根子都在计算图Computational Graph上。你不需要成为框架开发者但只要你在用PyTorch训练模型计算图就决定了你的显存怎么分配、梯度怎么流动、哪些操作能求导哪些不能、什么时候该用detach()、为什么有些in-place操作会报错。这些不是边角知识而是每天都会撞上的东西。这篇内容我想把计算图这件事从头讲透。不是那种计算图就是节点和边的教科书定义而是从实际写代码的角度出发动态图到底动态在哪、反向传播时图是怎么被消费的、叶子节点和非叶子节点的区别为什么会导致梯度拿不到、retain_graph什么时候必须开、torch.no_grad()和detach()到底差在哪。适合已经能跑通训练脚本、但遇到报错或显存问题只能靠搜索和试错的读者。看完之后你至少能做到看到报错信息就知道是哪类图的问题而不是盲目加torch.cuda.empty_cache()。2. 动态图的动态到底体现在哪里2.1 前向传播时图才被现场搭建PyTorch的计算图是动态的这句话很多人会背但真正理解它意味着什么的人不多。我用一个对比来说明。静态图框架比如早期的一些方案的思路是你先定义好整个计算流程框架把它编译成一张固定的图然后你往里面喂数据。图的结构在运行前就定死了循环、条件分支都得用框架提供的特殊语法来写。PyTorch完全不是这个路子。你写的是普通的Python代码for循环、if判断、甚至递归调用都随便用。每执行一行涉及张量的运算PyTorch就在背后记录一个节点。前向传播跑完图也就搭完了。下一次迭代代码重新执行一遍图重新搭一遍。这就是动态的含义——图的生命周期和一次前向传播绑定。这个特性带来的直接好处是调试极其方便。你可以在任意一行print张量的形状和值可以用Python的调试器单步跟踪可以随时改变网络结构。代价是每次迭代都有建图的开销不过对于绝大多数场景这点开销相比GPU计算可以忽略。我举个实际会遇到的例子。假设你要实现一个带early stopping的循环网络根据某个条件决定是否继续展开时间步def forward(self, x, max_steps10): h self.init_hidden(x) outputs [] for step in range(max_steps): h self.cell(x, h) outputs.append(h) if h.abs().mean() 1e-3: # 动态条件 break return torch.stack(outputs)这段代码在动态图框架里直接就能跑break什么时候触发取决于运行时的数据。如果换成静态图你得用框架提供的while_loop之类的原语来写麻烦得多。这就是动态图最直观的价值。2.2 每个张量身上的三个关键属性要理解计算图得先搞清楚一个张量身上挂了哪些和图相关的信息。我把它归纳成三个属性理解了这三个后面很多问题都能自己推导出来。第一个是requires_grad。这个布尔值决定这个张量是否需要计算梯度。默认情况下你手动创建的张量requires_gradFalse而模型参数nn.Parameter默认是True。只要一个运算的任意输入requires_gradTrue输出就会requires_gradTrue并且会在图中记录这个运算。第二个是grad_fn。这是指向创建该张量的函数的引用。如果你看到一个张量的grad_fn是None说明它是叶子节点不是由某个运算产生的。如果grad_fn有值比如AddBackward0说明它是由加法产生的反向传播时就会调用对应的反向函数。第三个是is_leaf。叶子节点是用户直接创建的、或者requires_gradFalse的张量。模型参数是典型的叶子节点。非叶子节点是由运算产生的中间结果。这个区分极其重要因为只有叶子节点的梯度会被默认保留在.grad属性里非叶子节点的梯度在反向传播后默认不保存。import torch x torch.tensor([1.0, 2.0], requires_gradTrue) # 叶子节点 y x * 2 # 非叶子节点 z y.sum() # 非叶子节点 print(x.is_leaf) # True print(y.is_leaf) # False print(y.grad_fn) # MulBackward0 object z.backward() print(x.grad) # tensor([2., 2.]) 叶子节点梯度保留 print(y.grad) # None 非叶子节点梯度默认不保留很多人第一次遇到为什么我拿不到中间层的梯度就是栽在这里。想拿非叶子节点的梯度得用y.retain_grad()显式告诉PyTorch保留它。2.3 图是用完即焚的动态图还有一个容易被忽略的特性反向传播完成后图默认就被释放了。这是为了省内存。你想想一个深层网络前向传播会产生大量中间激活值这些值在反向传播时都要用到所以必须存着。一旦反向传播算完这些中间值就没用了图也就没存在的必要了。这个设计导致一个常见报错你想对同一个loss做两次backward()第二次就会报Trying to backward through the graph a second time。因为第一次反向传播已经把图消费掉了。解决办法是第一次调用时加retain_graphTrue让图保留下来。但要注意保留图意味着中间激活值不释放显存占用会显著增加所以除非确实需要比如多任务loss分别反向否则不要随便开。我在实际项目里见过有人为了保险每次都加retain_graphTrue结果显存直接翻倍。这个参数不是用来防报错的它有明确的使用场景。3. 反向传播时图是怎么被消费的3.1 从loss到参数的梯度链路反向传播的本质是链式法则。计算图记录了前向传播的每一步运算反向传播就是从loss这个标量出发沿着图反向走一遍每经过一个节点就计算该节点的局部梯度然后乘上从上游传下来的梯度得到传给下游的梯度。用个简单的例子手动推一遍。假设x torch.tensor(3.0, requires_gradTrue) w torch.tensor(2.0, requires_gradTrue) b torch.tensor(1.0, requires_gradTrue) y w * x b # y 2*3 1 7 loss y ** 2 # loss 49前向传播建的图是x和w经过乘法得到wxwx和b经过加法得到yy经过平方得到loss。反向传播时loss y^2所以d(loss)/d(y) 2y 14y wx b所以d(y)/d(wx) 1d(y)/d(b) 1wx w * x所以d(wx)/d(w) x 3d(wx)/d(x) w 2链式相乘d(loss)/d(w) 14 * 1 * 3 42d(loss)/d(b) 14 * 1 14d(loss)/d(x) 14 * 1 * 2 28。PyTorch的grad_fn就是干这个的。每个grad_fn对象知道两件事怎么计算自己的局部梯度以及怎么把上游梯度传下去。backward()调用时从loss的grad_fn开始按拓扑逆序依次调用每个节点的反向函数。3.2 为什么梯度会累加而不是覆盖这是新手最容易踩的坑之一。PyTorch的梯度是累加的不是覆盖的。也就是说如果你连续做两次反向传播而没有清零梯度第二次的梯度会加到第一次上面。x torch.tensor([1.0], requires_gradTrue) y1 x * 2 y1.backward() print(x.grad) # tensor([2.]) y2 x * 3 y2.backward() print(x.grad) # tensor([5.]) 2 3不是3这个设计是有意为之的。在某些场景下比如RNN的梯度累积、多loss加权求和累加是需要的。但对于标准的训练循环你必须在每次反向传播前调用optimizer.zero_grad()否则梯度会越滚越大训练直接发散。我见过有人写训练循环时把zero_grad()放在backward()后面逻辑上好像也对——这次用完清掉下次用的时候是干净的。但这样写有个隐患如果某次迭代因为异常跳过了backward()梯度就没被清掉会污染下一次。标准写法是放在backward()之前每次迭代开头就清。3.3 梯度累积的实战用法既然梯度是累加的那就可以利用这个特性做梯度累积。当显存不够、没法开大batch时可以把一个大batch拆成几个小batch分别前向反向梯度自然累加最后再统一更新参数。效果等价于大batch训练。accumulation_steps 4 optimizer.zero_grad() for i, (inputs, labels) in enumerate(dataloader): outputs model(inputs) loss criterion(outputs, labels) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意loss要除以accumulation_steps否则累加后的梯度是原来的N倍相当于学习率放大了N倍。这个细节很多人会漏掉导致训练不稳定。提示梯度累积期间计算图在每个小batch反向传播后就被释放了所以显存占用和单个小batch相当不会因为累积而增加。这是它相比直接开大batch的优势。4. 那些和图相关的报错根因都在哪4.1 element 0 of tensors does not require grad这个报错的意思是你调用了backward()但loss这个张量的requires_grad是False没有图可以反向传播。常见原因有几个。第一种模型参数被冻结了。你可能手动把某些层的requires_grad设成了False或者用了torch.no_grad()包裹了前向传播。检查一下loss的requires_grad属性如果是False往上追溯是哪个环节断开了。第二种loss计算时用了不可导的操作。比如你把loss转成了Python float或者用了.item()或者经过了一个requires_gradFalse的张量运算。典型的是用accuracy之类的指标当lossaccuracy是离散的不可导。第三种输入数据没有requires_grad而且模型参数也被冻结。这种情况图根本建不起来。排查思路很简单从loss开始打印它的grad_fn和requires_grad然后往前找第一个grad_fn为None的地方那里就是断点。4.2 Trying to backward through the graph a second time前面提过图默认用完即焚。这个报错就是你试图对同一张图做第二次反向传播。除了加retain_graphTrue更值得思考的是你真的需要两次反向吗需要两次反向的典型场景是GAN。生成器和判别器的loss是分开的但都依赖于判别器对生成样本的输出。如果先算判别器loss反向图就被消费了再算生成器loss就没图了。这时候要么用retain_graphTrue要么把两个loss合在一起反向。# 方案一合并loss d_loss ... g_loss ... total_loss d_loss g_loss total_loss.backward() # 方案二retain_graph d_loss.backward(retain_graphTrue) g_loss.backward()方案一更省显存但两个loss的梯度会混在一起如果需要对不同参数做不同处理就不方便。方案二更灵活但显存占用高。实际用哪个取决于你的优化策略。还有一种情况是你在循环里反复对同一个loss反向这通常是代码逻辑写错了比如把backward()写在了内层循环里。这种要改代码结构而不是加retain_graph硬扛。4.3 in-place操作导致的版本冲突one of the variables needed for gradient computation has been modified by an inplace operation——这个报错也很常见。原因是反向传播需要用到前向传播时的某个中间值但你在前向之后、反向之前用in-place操作把这个值改了PyTorch检测到版本号不一致就报错了。x torch.tensor([1.0, 2.0], requires_gradTrue) y x * 2 y 1 # in-place操作y的版本号变了 z y.sum() z.backward() # 报错PyTorch用版本号机制来检测这种情况。每个张量有个_version属性in-place操作会递增它。反向传播时如果发现版本号对不上就知道中间值被篡改了。解决办法是用out-of-place操作替代y y 1 # 创建新张量不修改原张量或者用y y.clone() 1先克隆再操作。但克隆会增加显存所以优先用out-of-place写法。有些in-place操作是安全的比如relu_()在某些情况下不会触发报错因为它的反向不需要前向的输入值。但为了代码可维护性我建议除非有明确的性能需求否则统一用out-of-place写法。省那点显存不值得调试成本。4.4 显存泄漏的排查链路回到开头那个显存泄漏的例子。排查这类问题的完整链路是这样的第一步确认是不是计算图没释放。在训练循环里打印torch.cuda.memory_allocated()看它是否随step单调递增。如果是基本可以确定有东西没释放。第二步检查有没有把带梯度的张量存进了列表或字典。比如你把每个step的loss存进一个listloss是带图的图就一直挂着。正确做法是存loss.item()只存数值。第三步检查有没有在日志或调试代码里保存了中间张量。比如writer.add_scalar(grad, model.layer.weight.grad, step)如果weight.grad没被detach它可能带着图。第四步检查有没有忘记zero_grad()。虽然这不会导致图泄漏但会让梯度越来越大间接增加显存。第五步如果以上都排除了用torch.autograd.set_detect_anomaly(True)开启异常检测它会告诉你哪个操作导致了问题。这个模式很慢只在调试时开。我自己的经验是90%的显存泄漏都是把带梯度的张量存起来了这一类。养成习惯任何要长期保存的张量先.detach()或者.item()。5. no_grad、detach和inference_mode的取舍5.1 三者的语义差异这三个东西经常被混用但语义完全不同。torch.no_grad()是一个上下文管理器进入这个上下文后所有新创建的张量requires_grad都是False不会建图。它影响的是上下文内部的所有运算。detach()是张量的方法它返回一个新张量和原张量共享数据但requires_gradFalse且从图中脱离。原张量不受影响。torch.inference_mode()是后来引入的比no_grad更激进。它不仅不建图还会关闭一些autograd的元数据跟踪性能更好。但它有个限制在inference_mode下创建的张量之后不能在需要梯度的计算中使用。x torch.tensor([1.0], requires_gradTrue) with torch.no_grad(): y x * 2 print(y.requires_grad) # False z x * 2 z_detached z.detach() print(z.requires_grad) # True print(z_detached.requires_grad) # False5.2 验证和推理阶段该用哪个模型验证和推理时你不需要梯度用no_grad或inference_mode都能省显存、加速。选哪个如果你的代码只在推理阶段用且不涉及后续需要梯度的操作用inference_mode它更快。如果推理结果可能被用于后续需要梯度的计算比如某些元学习场景用no_grad更安全。我个人的习惯是纯推理用inference_mode验证阶段用no_grad。因为验证阶段有时会算一些指标虽然不需要梯度但代码可能和训练阶段共享用no_grad兼容性更好。注意no_grad和inference_mode都不会影响已经建好的图。如果你在上下文外面创建了带梯度的张量在上下文里面用它做运算运算结果不带梯度但原张量的图还在。要彻底断开得用detach()。5.3 detach的典型使用场景detach()最常用的场景是你想用某个网络的输出作为另一个网络的输入但不希望梯度回传到第一个网络。比如在强化学习里你用target network计算目标值这个目标值不应该对target network求梯度只对online network求梯度。这时候就要对target network的输出做detach()。with torch.no_grad(): target_q target_network(next_state) # 或者 target_q target_network(next_state).detach()两种写法效果类似但detach()更灵活因为它可以在图已经建好之后再断开。比如你先用target network算了一堆东西最后才决定哪部分要detach。还有一个场景是截断时间序列的反向传播。RNN展开太长时显存扛不住可以在某个时间步对hidden state做detach()让梯度只回传有限步。for step in range(seq_len): output, hidden rnn(input[step], hidden) if step % truncate_every 0: hidden hidden.detach()这样梯度最多回传truncate_every步显存可控。6. 自定义Function与图的手动干预6.1 什么时候需要自定义反向绝大多数情况下PyTorch内置的运算已经覆盖了你的需求。但有些场景需要自定义反向你实现了一个新的算子或者你想在前向和反向之间做一些特殊处理或者你想用更高效的数值稳定实现。自定义Function需要继承torch.autograd.Function实现forward和backward两个静态方法。forward里做计算backward里根据上游梯度计算对输入的梯度。class MyReLU(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x.clamp(min0) staticmethod def backward(ctx, grad_output): x, ctx.saved_tensors grad_input grad_output.clone() grad_input[x 0] 0 return grad_input关键点是ctx.save_for_backward它把反向需要的前向中间值存起来。注意只能存张量存其他类型用ctx.other value。存张量必须用save_for_backward否则可能出问题。6.2 自定义Function的常见坑第一个坑是backward的返回值数量必须和forward的输入数量一致。如果某个输入不需要梯度返回None。第二个坑是backward里不要做in-place操作修改grad_output因为它可能被其他地方引用。要改先clone()。第三个坑是save_for_backward存的张量如果在前向之后被修改了反向时会报版本错误。所以存进去之后就别动它。第四个坑是forward里如果用了no_grad上下文save_for_backward仍然有效但要注意上下文的范围。我实际用自定义Function最多的场景是实现一些数值不稳定的loss。比如某些涉及log-sum-exp的操作直接用内置函数可能溢出自己写一个数值稳定的版本同时自定义反向保证梯度正确。6.3 用hook观察和干预梯度除了自定义FunctionPyTorch还提供了hook机制可以在不修改模型代码的情况下观察或修改梯度。register_hook可以注册在张量上反向传播经过这个张量时触发def print_grad(grad): print(grad.abs().mean()) x torch.tensor([1.0, 2.0], requires_gradTrue) y x * 2 y.register_hook(print_grad) z y.sum() z.backward()register_full_backward_hook可以注册在模块上观察模块输入输出的梯度。这在调试梯度消失或爆炸时很有用。for name, module in model.named_modules(): module.register_full_backward_hook( lambda m, gin, gout: print(f{name}: {gout[0].abs().mean()}) )hook的返回值如果非None会替换原来的梯度。这个特性可以用来做梯度裁剪、梯度缩放等操作。但要注意hook里修改梯度是in-place的可能影响其他hook。7. 计算图视角下的显存优化实践7.1 激活值重计算反向传播需要前向的中间激活值这些值占显存的大头。激活值重计算也叫gradient checkpointing的思路是前向传播时只存少数几个关键节点反向传播需要中间值时重新算一遍。用计算换显存。PyTorch提供了torch.utils.checkpointfrom torch.utils.checkpoint import checkpoint def forward(self, x): x checkpoint(self.block1, x) x checkpoint(self.block2, x) return x每个checkpoint包裹的模块前向时不存中间激活反向时重新前向一遍。显存能省很多代价是训练速度慢一些大概多一次前向的时间。用checkpoint有几个注意点。第一被包裹的模块输入必须requires_gradTrue否则checkpoint不生效。第二如果模块里有随机性比如dropout前向重算时结果可能不一致需要用preserve_rng_stateTrue。第三checkpoint对batch norm不太友好因为重算时统计量可能变了。7.2 混合精度与图的关系混合精度训练AMP用float16做前向和反向float32做参数更新。它和计算图的关系在于autograd在float16下也能正常工作但梯度可能下溢。AMP的GradScaler就是用来解决这个问题的——它把loss放大反向得到的梯度也放大更新前再缩回去。scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意scaler.scale(loss)返回的是一个新的loss张量它带着缩放系数。反向传播时梯度也被缩放scaler.step会先检查有没有inf/nan没有才更新参数。AMP下计算图的结构和float32一样只是数据类型变了。但要注意有些操作在float16下不稳定比如softmax、layer normautocast会自动把这些操作转成float32。你不需要手动干预但要知道这个机制存在。7.3 梯度检查点与no_grad的配合有个容易混淆的点checkpoint内部的前向重算是在no_grad下做的吗不是。重算时需要建图否则反向没法继续。checkpoint的实现是第一次前向在no_grad下做不存图反向时重新前向这次建图然后立即反向用完释放。所以checkpoint包裹的模块其内部不能有依赖全局状态的运算因为重算时状态可能变了。这也是为什么checkpoint对dropout、batchnorm要特别处理。我实际用checkpoint的经验是在Transformer类模型上效果最明显因为每层都有大量激活值。一般从中间几层开始加checkpoint前面几层激活值小加了收益不大后面几层如果接近输出也可以不加。具体加在哪用显存分析工具看哪层占用大就加哪层。8. 一些容易被忽略的边界情况8.1 叶子节点的梯度什么时候是None叶子节点默认保留梯度但有些情况下grad是None。第一种这个叶子节点没有参与loss计算。比如你定义了两个参数但loss只用了其中一个另一个的grad就是None。第二种这个叶子节点的requires_gradFalse。第三种反向传播还没执行。optimizer.step()会跳过grad为None的参数所以不会报错。但如果你手动更新参数要检查grad是否为None。8.2 高阶导数的图PyTorch支持高阶导数但默认只建一阶图。要算二阶导需要在反向传播时也建图x torch.tensor([2.0], requires_gradTrue) y x ** 3 dy torch.autograd.grad(y, x, create_graphTrue)[0] d2y torch.autograd.grad(dy, x)[0] print(d2y) # 6*x 12create_graphTrue让反向传播的过程也被记录从而可以再反向。高阶导数在元学习、物理模拟等场景有用但显存开销大因为图会嵌套。8.3 多GPU下的图DataParallel和DistributedDataParallel下每个GPU上的前向传播建各自的图反向传播时梯度汇总到主GPU。图本身不跨GPU每个GPU独立管理自己的图。所以显存优化要针对单卡来做多卡只是把batch拆开了。DDP比DP效率高的原因之一就是DDP在每个进程里独立建图通信只在梯度汇总时发生而DP有额外的图复制开销。8.4 图的可视化想直观看看图长什么样可以用torchviz或tensorboard的graph功能。torchviz最简单from torchviz import make_dot y model(x) make_dot(y, paramsdict(model.named_parameters())).render(graph, formatpng)生成的图会显示每个运算节点和梯度函数。对于复杂模型图会很大但看看关键路径还是有帮助的。我一般只在调试奇怪的反向问题时才用平时靠grad_fn和报错信息就够了。9. 我踩过的几个真实坑第一个坑是关于loss.item()的。早期我写日志时直接存loss没加.item()结果每个step的loss都带着图跑了几百步显存就爆了。后来改成loss.item()问题消失。这个坑的本质是Python的list会持有对象的引用带图的张量被list引用着图就不会释放。第二个坑是关于retain_graph的。有次做多任务学习两个loss分别反向我加了retain_graphTrue结果显存翻倍。后来改成把两个loss加起来一次反向显存恢复正常。多任务loss如果不需要分别处理梯度合并反向是最优解。第三个坑是关于in-place的。我在一个自定义模块里用了x residual训练时报版本错误。改成x x residual就好了。但这样会多一份显存对于大模型有影响。后来我改成在残差连接处用x x residual其他地方尽量用in-place平衡了显存和正确性。第四个坑是关于detach()的。有次做target network更新我忘了对target的输出做detach()结果梯度回传到了target networktarget network被意外更新了训练完全跑偏。这个bug很难发现因为loss看起来在降但效果就是不对。后来我养成了习惯任何不应该被更新的网络输出一律detach()。第五个坑是关于no_grad的。有次在验证阶段用了no_grad但验证代码里有一行把验证loss存进了全局变量训练阶段又用这个loss做了一些计算。因为验证loss是在no_grad下算的没有图训练阶段用的时候报错。这个问题的根源是训练和验证的代码耦合太紧后来我把它们彻底分开各自管理自己的计算图。这些坑的共同点是它们都不是PyTorch的bug而是对计算图生命周期理解不到位导致的。理解了图什么时候建、什么时候释放、什么时候断开这些问题都能自己推导出答案。计算图这个东西刚开始学的时候觉得抽象但用久了会发现它其实是PyTorch里最讲道理的部分。每一个报错都有明确的物理含义每一个参数都有清晰的使用场景。你不需要背下来所有细节但建立起图的生命周期这个心智模型之后遇到问题就知道往哪个方向查。我现在的习惯是看到任何autograd相关的报错先问自己三个问题图建起来了吗图被消费了吗图被意外保留了吗这三个问题能解决大部分问题。