企业官网建设流程全解析

1. 训练到第N个epoch突然OOM：先别急着调batch size

显存泄漏这个问题，我敢说每个用PyTorch跑过稍大规模训练的人都遇到过。现象特别典型：前几个epoch显存占用稳稳的，跑到中途nvidia-smi一看，显存曲线跟爬楼梯似的，一个epoch比一个epoch高，最后在某个随机的时间点直接甩给你一个CUDA out of memory。很多人的第一反应是把batch size砍一半，或者上torch.cuda.empty_cache()，结果发现要么治标不治本，要么训练速度掉得没法看。

这里要先厘清一个概念：显存泄漏和显存峰值过高是两码事。峰值过高是模型本身或某个操作瞬时吃掉了大量显存，比如attention矩阵在序列长度拉长时平方级膨胀；而泄漏是显存占用随着训练步数单调递增，即使输入尺寸完全不变。判断方法很简单——固定输入shape，跑几百个step，每隔几十步打印一次torch.cuda.memory_allocated()和torch.cuda.memory_reserved()，如果allocated持续上涨，那就是实打实的泄漏。

这篇文章面向的是已经能跑通训练脚本、但被显存问题卡住的开发者。我会从计算图的持有关系讲起，把常见的泄漏源头一个个拆开，给出可复现的排查链路和修复代码。所有方案都在PyTorch 1.11到2.x上验证过，不依赖任何特殊硬件。

提示：排查前先确认你的环境里torch.cuda.is_available()为True，且驱动版本与CUDA版本匹配。环境本身有问题的话，后面所有分析都是白费功夫。

2. 计算图为什么会被"扣住"：从autograd的引用链说起

2.1 一张计算图的完整生命周期

PyTorch的动态图机制决定了每次前向传播都会构建一张新的计算图。这张图由Function节点和Tensor边组成，每个非叶子节点的Tensor都持有grad_fn，而grad_fn又通过next_functions指向它的输入节点。正常情况下，你调用loss.backward()之后，PyTorch会释放掉这张图里保存的中间激活值，只保留叶子节点的.grad。

问题出在"正常情况下"这四个字。只要有任何一条Python引用链还指向计算图中的某个非叶子Tensor，整张图就无法被释放。这就像你搬家时把钥匙落在了旧房子里，只要钥匙还在，搬家公司就没法把房子清空。常见的"钥匙"包括：被append进list的loss、被存进dict的中间输出、被闭包捕获的Tensor、被全局变量引用的模型输出。

2.2 一个最小可复现的泄漏demo

先看一段能稳定复现泄漏的代码，后面所有分析都围绕它展开：

import torch import torch.nn as nn class Net(nn.Module): def __init__(self): super().__init__() self.fc = nn.Linear(1024, 1024) def forward(self, x): return self.fc(x) model = Net().cuda() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) losses = [] # 罪魁祸首 for step in range(1000): x = torch.randn(64, 1024).cuda() y = model(x) loss = y.sum() losses.append(loss) # 持有计算图引用 loss.backward() optimizer.step() optimizer.zero_grad() if step % 100 == 0: print(step, torch.cuda.memory_allocated() / 1024**2, "MB")

跑一下你会发现，显存占用随着step线性增长。原因就是losses这个list把每个step的loss都存了下来，而loss是计算图的根节点，持有它就等于持有整张图。修复方式有两种：存loss.item()（转成Python float，脱离计算图），或者存loss.detach()（共享存储但切断梯度链）。

2.3 detach和item的本质区别

很多人分不清detach()和item()。detach()返回的是一个与原Tensor共享底层存储、但requires_grad=False的新Tensor，它仍然占用显存；item()则把标量Tensor转成Python的float对象，彻底脱离GPU。所以如果你只是想记录loss曲线，用item()；如果后续还要对这个值做GPU上的运算，用detach()。

注意：with torch.no_grad()包裹的代码块里产生的Tensor默认不带梯度，但如果你在块内把Tensor赋给了块外的变量，它依然会存活。no_grad只影响新创建的Tensor，不影响已有引用。

3. 五类高频泄漏源：从loss累积到自定义autograd

3.1 累积loss、metric和中间输出

这是最常见的一类。除了上面demo里的list，还有几种变体：把每个step的output存进list做可视化、把hidden_state缓存起来做时序分析、在验证循环里忘了切eval()模式导致BN层持续更新running stats并保留图。排查方法是在训练循环里加一行gc.get_referrers(loss)，看看除了局部变量还有谁在引用它。

修复原则很简单：任何跨越step边界存活的对象，都必须脱离计算图。记录指标用.item()，缓存特征用.detach().cpu()，需要保留梯度的场景（比如梯度累积）要显式管理生命周期。

3.2 优化器状态与参数注册的隐性持有

优化器本身会持有参数的引用，这是正常的。但如果你在训练过程中动态地往model里注册新的buffer或parameter，比如self.register_buffer('running_feat', feat)，而这个feat是从计算图里来的，那优化器或模型就会一直扣着这张图。还有一种情况是用了torch.nn.ParameterList或ModuleList，在forward里动态append，导致每次前向都新增参数，显存自然只增不减。

检查方法：训练几个step后打印len(list(model.parameters()))和len(list(model.buffers()))，如果数字在涨，那就是动态注册的问题。

3.3 自定义autograd.Function的ctx泄漏

写自定义算子时，forward里的ctx.save_for_backward()如果保存了不该保存的Tensor，或者backward里返回的梯度引用了外部变量，都会造成泄漏。更隐蔽的是在forward里把中间结果存到了self上：

class MyFunc(torch.autograd.Function): @staticmethod def forward(ctx, x): ctx.save_for_backward(x) MyFunc.cache = x * 2 # 类属性持有，永久泄漏 return x * 2

这种类属性或模块级变量持有Tensor的写法，是很多"跑着跑着就OOM"的元凶。修复方式是改用ctx.save_for_backward，并确保backward返回的梯度不引用任何外部Tensor。

3.4 DataLoader的worker与pin_memory

DataLoader的num_workers > 0时，每个worker进程会独立加载数据。如果collate_fn里返回了CUDA Tensor，或者dataset的__getitem__里做了GPU操作，worker进程会持有显存。另外pin_memory=True会使用锁页内存，虽然不直接吃显存，但配合worker的预取队列（prefetch_factor）会放大内存占用。

排查时把num_workers设为0跑一遍，如果泄漏消失，问题就在数据加载侧。修复方案是确保dataset只返回CPU Tensor，把.cuda()操作放到训练循环里。

3.5 全局缓存与闭包捕获

最后这类最阴间：在Jupyter Notebook里反复运行同一个cell，旧的模型、优化器、Tensor都还挂在_变量或Out字典里；或者用functools.lru_cache缓存了带Tensor的函数返回值；再或者用闭包捕获了训练循环里的变量。这类泄漏在脚本模式下不容易出现，但在交互式环境里极其常见。

4. 定位工具链：memory_summary、snapshot与引用追踪

4.1 torch.cuda.memory_summary的读法

PyTorch内置的torch.cuda.memory_summary()能打印出当前显存的详细分布，包括allocated、reserved、active、inactive等指标。关键看两个数：allocated是实际被Tensor占用的显存，reserved是PyTorch向CUDA申请的总量。如果allocated稳定但reserved持续涨，说明是碎片化问题，不是泄漏；如果allocated本身在涨，那才是真泄漏。

print(torch.cuda.memory_summary(device=None, abbreviated=False))

输出里会按Tensor大小分桶统计，能快速看出是大Tensor泄漏还是小Tensor堆积。

4.2 用memory_snapshot做逐行归因

torch.cuda.memory_snapshot()返回一个列表，每个元素代表一个内存块，包含size、allocated_size、frames等字段。frames里记录了分配这个块时的Python调用栈。你可以写个小脚本，把snapshot按frames聚合，找出哪个文件哪一行分配的显存最多：

import torch, collections snap = torch.cuda.memory_snapshot() counter = collections.Counter() for block in snap: if block['allocated_size'] > 0: for frame in block['frames']: counter[frame['filename'] + ':' + str(frame['line'])] += block['allocated_size'] for loc, size in counter.most_common(10): print(loc, size / 1024**2, "MB")

这个方法的威力在于，它能直接告诉你泄漏发生在你的代码的哪一行，而不是笼统地说"某个Tensor没释放"。

4.3 gc与weakref追踪引用链

当snapshot定位到具体行之后，用gc.get_referrers()反查谁在引用这个对象：

import gc target = some_leaked_tensor for ref in gc.get_referrers(target): print(type(ref), ref)

如果引用者是list、dict或某个模块的属性，基本就能锁定泄漏源。对于复杂场景，可以用weakref.ref给可疑对象加弱引用回调，在对象被回收时打印日志，确认它到底有没有被释放。

提示：memory_snapshot在PyTorch 1.10以上才比较完善，老版本建议升级。另外snapshot本身会消耗一定显存，别在OOM边缘调用。

5. 修复实战：从empty_cache的误用到梯度累积的正确姿势

5.1 empty_cache到底该不该用

torch.cuda.empty_cache()的作用是把PyTorch reserved但未使用的显存还给CUDA驱动，让其他进程能用。它不解决泄漏，只是把已经泄漏的显存从"reserved"变成"free"，但那些被引用的Tensor依然占着allocated。频繁调用还会导致后续分配变慢，因为PyTorch得重新向驱动申请。

正确用法：在显存碎片化严重、且确认没有泄漏的前提下，偶尔调用一次。比如在验证循环开始前，或者模型切换时。把它当成"整理房间"，而不是"扔垃圾"。

5.2 梯度累积的显存陷阱

梯度累积的常见写法是：

for i, (x, y) in enumerate(loader): loss = model(x, y) / accum_steps loss.backward() if (i + 1) % accum_steps == 0: optimizer.step() optimizer.zero_grad()

这里如果loss被存进了list，或者model(x, y)的输出被缓存，泄漏照样发生。更隐蔽的是，有些人为了"节省显存"在累积期间不调用zero_grad()，结果梯度一直累加，虽然不算泄漏，但会导致梯度爆炸。梯度累积的正确姿势是：每次backward后图自动释放，只要不持有loss引用，显存就是稳定的。

5.3 用del和gc手动干预的时机

在确认某个大Tensor不再需要时，可以显式del tensor并调用gc.collect()。但要注意，del只是删除引用，如果还有其他引用存在，对象依然存活。在训练循环里，局部变量会在下一次迭代时被覆盖，通常不需要手动del。真正需要手动干预的场景是：在循环外创建了大Tensor，循环内反复使用，用完后想提前释放。

big_cache = torch.randn(10000, 10000).cuda() for step in range(1000): out = model(big_cache) # ... 训练逻辑 del big_cache torch.cuda.empty_cache()

5.4 一个完整的修复案例

回到第2节的demo，完整修复版本如下：

import torch import torch.nn as nn class Net(nn.Module): def __init__(self): super().__init__() self.fc = nn.Linear(1024, 1024) def forward(self, x): return self.fc(x) model = Net().cuda() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) losses = [] for step in range(1000): x = torch.randn(64, 1024).cuda() y = model(x) loss = y.sum() losses.append(loss.item()) # 关键修复：脱离计算图 loss.backward() optimizer.step() optimizer.zero_grad() if step % 100 == 0: print(step, torch.cuda.memory_allocated() / 1024**2, "MB")

改动只有一处，但显存曲线会从爬楼梯变成一条水平线。这就是定位到根因之后修复的威力——不需要调batch size，不需要加empty_cache，一行代码解决。

6. 几个反直觉的坑：no_grad、inference_mode与BN的隐藏开销

6.1 no_grad不等于不占显存

torch.no_grad()只是让新创建的Tensor不带梯度，但前向传播的中间激活值依然会占用显存，只是backward时不需要它们，所以可以在前向结束后立即释放。如果你在no_grad块里把中间结果存到了外部变量，显存照样涨。验证循环里最常见的错误就是：

with torch.no_grad(): for x, y in val_loader: out = model(x) all_outs.append(out) # 依然持有Tensor

正确做法是all_outs.append(out.cpu())或者只记录指标。

6.2 inference_mode比no_grad更彻底

PyTorch 1.9引入了torch.inference_mode()，它比no_grad更激进：不仅不记录梯度，还会禁用版本计数和autograd相关的元数据，显存占用和速度都更优。纯推理场景建议优先用它。但注意，inference_mode下创建的Tensor不能参与后续需要梯度的计算，否则会报错。

6.3 BatchNorm的running stats与显存

BN层在训练模式下会更新running_mean和running_var，这些是buffer，不占多少显存。但如果你的模型在验证时忘了model.eval()，BN会继续用batch统计量并更新running stats，同时因为处于训练模式，前向会构建计算图，显存占用直接翻倍。这是一个非常隐蔽的坑，尤其在自定义训练循环里容易漏掉。

6.4 多卡训练下的显存不均

用DataParallel时，显存占用往往集中在主卡上，因为loss和梯度汇总都在主卡。DistributedDataParallel则相对均衡。如果你发现只有某张卡OOM，先确认是不是用了DataParallel。另外，多卡场景下每张卡都会维护一份完整的模型副本和优化器状态，显存需求是单卡的N倍（N为卡数），这个账要提前算清楚。

7. 把显存监控做成习惯：训练脚本里的自检清单

7.1 每个epoch打印一次显存指标

在训练循环里加几行监控代码，成本极低，收益极高：

def log_memory(step): alloc = torch.cuda.memory_allocated() / 1024**2 reserv = torch.cuda.memory_reserved() / 1024**2 print(f"step {step}: allocated={alloc:.1f}MB reserved={reserv:.1f}MB")

如果allocated在固定输入下持续上涨，立刻停下来排查，别等到OOM才动手。

7.2 用tracemalloc配合CPU侧排查

显存泄漏有时候是CPU侧对象泄漏的连带结果。tracemalloc能追踪Python对象的内存分配，配合gc可以找出哪些Python对象在增长。虽然它不直接追踪GPU显存，但能帮你发现list、dict的无界增长。

7.3 训练前的显存预算

跑之前先算一笔账：模型参数显存 = 参数量 × 4字节（fp32）；优化器状态（以Adam为例）= 参数量 × 8字节（一阶矩+二阶矩）；激活值显存与batch size、序列长度、层数成正比。把这三部分加起来，再留20%余量，就是你的显存下限。如果预算超了，先调模型结构或batch size，别指望靠empty_cache硬撑。

7.4 一个可复用的排查流程

我把上面所有内容浓缩成一个排查流程，遇到OOM时按顺序走：

1. 固定输入shape，打印allocated曲线，确认是否泄漏。
2. 检查所有跨step存活的容器（list、dict、全局变量），确保存的是item或detach后的CPU Tensor。
3. 把num_workers设为0，排除DataLoader侧问题。
4. 用memory_snapshot定位到具体代码行。
5. 用gc.get_referrers反查引用链。
6. 修复后重跑，确认allocated曲线水平。

这套流程我在多个项目上用过，从简单的图像分类到复杂的多模态训练，基本都能在半小时内定位到根因。显存泄漏不可怕，可怕的是盲目调参和滥用empty_cache。把计算图的引用关系搞清楚，问题就解决了一大半。