1. 显存泄漏到底是个什么问题
训练脚本跑着跑着，nvidia-smi里的显存占用像爬楼梯一样一路往上走，从最初的 8GB 慢慢涨到 10GB、12GB，最后在某个 step 直接甩出一个CUDA out of memory，进程崩掉。更气人的是，你把 batch size 调小、把图片尺寸降下来，它还是会在差不多的时间点崩——因为根因根本不是模型太大，而是显存泄漏。
显存泄漏和内存泄漏在表象上很像，但排查手段完全不同。CPU 内存泄漏可以用tracemalloc、objgraph这类工具去追对象引用链，而显存泄漏发生在 GPU 上，Python 的 GC 管不到 CUDA 缓存分配器那一层，你得用torch.cuda.memory_allocated()、torch.cuda.memory_reserved()这些接口去观测，再配合gc和引用计数去定位。
这篇文章面向的是已经能跑通训练、但被显存持续增长折磨的从业者。我会把显存泄漏的常见成因、定位手法、修复方案按实战顺序讲清楚，包括计算图没释放、张量被意外持有、CUDA 缓存碎片、DataLoader worker 泄漏这几类高频问题。文中给出的代码片段都可以直接抄进你的训练脚本里用。
先说一个基本判断：显存占用持续单调上升，且不随 step 回落，基本可以判定为泄漏。正常的训练显存曲线应该是锯齿状的——前向时涨、反向释放后回落，在一个区间内波动。如果你的曲线是阶梯状单调上升，那就要开始查了。
2. 先搞清楚 PyTorch 的显存账本
2.1 allocated 与 reserved 的区别
很多人一上来就盯着nvidia-smi看，其实那个数字包含了 CUDA context、cuDNN 工作区、其他进程占用等一堆东西，不够精确。PyTorch 自己提供了更细的观测口径：
torch.cuda.memory_allocated()：当前真正被张量占用的显存，单位字节。这是你的模型参数、激活值、梯度实际占的量。torch.cuda.memory_reserved()：PyTorch 向 CUDA 申请并缓存起来的显存。缓存分配器为了减少cudaMalloc调用，会把释放的张量内存留在池子里复用，所以 reserved 通常大于等于 allocated。torch.cuda.max_memory_allocated()：历史峰值，用来判断峰值是否异常。
关键点在于：allocated 持续上涨才是真泄漏。如果只是 reserved 涨、allocated 稳定，那多半是缓存碎片问题，可以通过torch.cuda.empty_cache()缓解，但注意这个操作会拖慢训练，不建议在训练循环里频繁调用。
2.2 一个最小可用的显存监控函数
我在每个项目里都会塞一个这样的工具函数，训练循环里每隔 N 步打印一次：
import torch def log_gpu_memory(step, tag=""): allocated = torch.cuda.memory_allocated() / 1024**2 reserved = torch.cuda.memory_reserved() / 1024**2 max_allocated = torch.cuda.max_memory_allocated() / 1024**2 print(f"[step {step}] {tag} " f"allocated={allocated:.1f}MB " f"reserved={reserved:.1f}MB " f"max={max_allocated:.1f}MB")把它插在optimizer.step()之后、loss.backward()之前各打一次，对比两个位置的 allocated 差值。如果每个 step 的差值都在缓慢累积，泄漏就坐实了。
提示：
torch.cuda.reset_peak_memory_stats()可以在你想重新统计峰值时调用，比如每个 epoch 开始时重置一次，方便观察单 epoch 内的峰值变化。
2.3 为什么 reserved 会一直涨
CUDA 缓存分配器按块管理显存，块大小有 512B、1KB、2MB 等多档。当你的张量尺寸变化频繁（比如变长序列、动态 shape），分配器会产生大量不同尺寸的块，释放后这些块因为尺寸不匹配无法被后续请求复用，就形成了碎片。碎片多了，reserved 就会虚高，最终触发 OOM。
这种情况的典型特征是：allocated 稳定，但 reserved 缓慢上涨，且empty_cache()能明显把 reserved 降下来。解决办法是尽量固定输入 shape（padding 到固定长度）、避免在训练循环里创建临时大张量。
3. 计算图没释放：最常见的泄漏源头
3.1 计算图持有机制回顾
PyTorch 的动态图在loss.backward()之后，理论上计算图会被释放。但有几个常见操作会让图一直被持有：
- 把
loss存进了列表或字典，比如losses.append(loss)用于后续统计。这个loss是带 grad_fn 的张量，它持有整条计算图的引用。 - 在
backward()之前做了loss.item()之外的张量操作，比如total_loss += loss（而不是total_loss += loss.item()）。 - 用了
retain_graph=True但没在合适时机释放。
我见过最典型的写法是这样的：
running_loss = 0.0 for step, (x, y) in enumerate(loader): out = model(x) loss = criterion(out, y) loss.backward() optimizer.step() optimizer.zero_grad() running_loss += loss # 错误：loss 是张量，持有计算图running_loss从 float 变成张量后，每个 step 的 loss 计算图都被累加进去，显存自然一路涨。正确写法是running_loss += loss.item()，.item()会把张量转成 Python float，切断计算图引用。
3.2 用 gc 和引用计数定位
如果你不确定是哪个变量在持有计算图，可以用gc.get_referrers()反查。下面这段代码能帮你找到持有某个张量的对象：
import gc import torch def find_tensor_referrers(tensor, topk=5): referrers = gc.get_referrers(tensor) for i, ref in enumerate(referrers[:topk]): print(f"--- referrer {i} ---") print(type(ref)) if isinstance(ref, dict): for k, v in list(ref.items())[:5]: print(f" key={k}, type={type(v)}") elif isinstance(ref, (list, tuple)): print(f" len={len(ref)}")在怀疑泄漏的 step 上，拿到一个带 grad_fn 的张量，调这个函数看看谁在引用它。实测下来，十有八九能揪出那个不该存在的 list 或 dict。
3.3 修复清单
针对计算图泄漏，我整理了一份检查清单，按优先级排序：
| 检查项 | 错误写法 | 正确写法 |
|---|---|---|
| loss 累加 | total += loss | total += loss.item() |
| 保存中间张量 | feats.append(feat) | feats.append(feat.detach()) |
| 日志记录 | writer.add_scalar('loss', loss) | writer.add_scalar('loss', loss.item()) |
| 早停判断 | if loss < best: ... | if loss.item() < best: ... |
| 梯度裁剪 | 在 backward 前调用 | 在 backward 后、step 前调用 |
注意：
detach()只是切断计算图，张量本身还在显存里。如果feats列表无限增长，即使 detach 了也会泄漏。这种情况要么限制列表长度，要么定期清空。
4. 张量被意外持有：那些看不见的引用
4.1 全局变量和闭包
Python 的闭包和全局变量是显存泄漏的隐形杀手。比如你在函数里定义了一个 hook，hook 里引用了某个中间张量，而这个 hook 又注册在全局的 module 上，那这个张量就永远不会被释放。
activations = {} def get_activation(name): def hook(model, input, output): activations[name] = output # output 持有计算图 return hook model.layer1.register_forward_hook(get_activation('layer1'))这段代码在推理时没问题，但训练时activations会一直累积，而且每个 output 都带 grad_fn。修复方式是activations[name] = output.detach()，或者用完就activations.clear()。
4.2 DataLoader 的 worker 泄漏
num_workers > 0时，每个 worker 是独立进程，它们会预取 batch。如果 worker 里的 dataset 持有大张量（比如把整个数据集加载到 GPU 上），或者collate_fn返回了带计算图的张量，就会造成泄漏。
一个隐蔽的坑是：在 dataset 的__getitem__里做了 GPU 操作。worker 进程默认不共享 CUDA context，你在 worker 里创建 GPU 张量，主进程拿到的可能是副本，而且 worker 退出时这些张量不一定被正确释放。
class MyDataset(Dataset): def __getitem__(self, idx): img = self.images[idx] # 错误：在 worker 里做 GPU 操作 # img = img.cuda() return img正确做法是 dataset 只做 CPU 侧的加载和预处理，.cuda()放到训练循环里做。另外pin_memory=True配合non_blocking=True能加速传输，但要注意 pin memory 本身也占内存。
4.3 优化器状态和 EMA
Adam 优化器会为每个参数维护 exp_avg 和 exp_avg_sq 两个状态，显存占用是参数量的两倍。如果你在训练中途重建了优化器但没释放旧的，或者用了 EMA（指数移动平均）却忘了 detach，都会导致显存增长。
EMA 的正确写法：
class EMA: def __init__(self, model, decay=0.999): self.decay = decay self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items()} @torch.no_grad() def update(self, model): for k, v in model.state_dict().items(): self.shadow[k] = self.decay * self.shadow[k] + (1 - self.decay) * v关键是@torch.no_grad()和detach()，否则 shadow 里的张量会带着计算图，越积越多。
5. 实操：一步步定位泄漏点
5.1 二分法缩小范围
不要一上来就读代码找问题，先用二分法把泄漏范围缩小。具体做法是：
- 写一个最小复现脚本，只保留模型、数据加载、训练循环三部分。
- 先跑 10 个 step，记录 allocated 变化。
- 如果泄漏，注释掉
loss.backward()，只跑前向，看是否还泄漏。 - 如果前向不泄漏，加上 backward，再测。
- 逐步加回 optimizer.step()、日志、验证等模块，定位到具体环节。
这个方法我用了很多次，通常半小时内能把范围缩到几十行代码。
5.2 用 memory snapshot 抓现场
PyTorch 提供了torch.cuda.memory._record_memory_history()和_dump_snapshot()，能记录每个显存分配的调用栈。这是定位泄漏的核武器：
import torch # 开始记录 torch.cuda.memory._record_memory_history(max_entries=100000) # 跑一段训练 for step, batch in enumerate(loader): train_step(batch) if step == 50: break # 导出快照 torch.cuda.memory._dump_snapshot("mem_snapshot.pickle") torch.cuda.memory._record_memory_history(enabled=None)导出的 pickle 文件可以上传到 PyTorch 官方的 memory visualizer 工具里查看（搜索 "pytorch memory visualizer" 就能找到在线版）。它会画出显存随时间的变化曲线，以及每个分配点的调用栈，泄漏点一目了然。
5.3 一个真实的排查案例
我之前遇到过一个 mask2former 训练显存持续增长的问题。用 snapshot 一看，发现是Matcher模块里的cost_matrix被缓存了。原因是代码里写了一个@lru_cache装饰器缓存匹配结果，但 cache 的 key 包含了张量，导致每次迭代都新增一条缓存记录，且旧记录不释放。
修复方式是把lru_cache的maxsize设成一个小值，或者干脆去掉缓存。这个坑很隐蔽，因为lru_cache在 CPU 侧看起来人畜无害，但它缓存的张量会一直占着显存。
提示：任何带缓存的装饰器（
lru_cache、cache、自定义 memoize）用在涉及张量的函数上时，都要格外小心。缓存 key 如果是张量，比较的是对象 id 而非值，会导致缓存无限增长。
6. 常见问题速查与避坑技巧
6.1 问题速查表
| 现象 | 可能原因 | 排查方法 | 修复方案 |
|---|---|---|---|
| allocated 单调上升 | 计算图未释放 | 检查 loss 累加、日志记录 | 用.item()、.detach() |
| reserved 高但 allocated 稳 | 缓存碎片 | empty_cache()后观察 | 固定输入 shape |
| 训练几个 epoch 后 OOM | 验证集未no_grad | 检查 val 循环 | 加@torch.no_grad() |
| 多卡训练显存不均 | 梯度同步问题 | 看各卡 allocated | 检查 DDP 配置 |
| DataLoader 报错后显存不释放 | worker 进程残留 | ps aux看进程 | 设persistent_workers=False |
| 中途重建模型后显存翻倍 | 旧模型未释放 | del model; gc.collect() | 显式删除并清缓存 |
6.2 验证循环的经典坑
验证循环忘记加torch.no_grad()是最常见的低级错误，但它的表现很迷惑：训练时显存正常，一到验证就涨，验证完也不回落。因为验证时的前向会构建计算图，而你没有 backward，这些图就一直挂着。
# 错误 def validate(model, loader): model.eval() for x, y in loader: out = model(x) loss = criterion(out, y) # 没有 no_grad，计算图一直累积 # 正确 @torch.no_grad() def validate(model, loader): model.eval() for x, y in loader: out = model(x) loss = criterion(out, y)@torch.no_grad()装饰器比在函数内写with torch.no_grad():更不容易漏，推荐用装饰器。
6.3 几个我踩过的坑
第一个坑是torch.cuda.empty_cache()的滥用。有人把它放进训练循环每个 step 都调，结果训练速度掉了一半。这个函数会释放所有未使用的缓存块，下次分配又要重新cudaMalloc，开销很大。正确用法是在 epoch 之间调，或者显存紧张时手动调一次。
第二个坑是pin_memory配合大 batch。pin memory 会把数据锁在页锁定内存里，如果 batch 很大且num_workers很多，主机内存会先爆掉。这时候报的错是主机 OOM，不是显存 OOM，容易误判。
第三个坑是混合精度训练里的GradScaler。如果scaler.step(optimizer)之后忘了scaler.update()，或者scaler被重复创建，会导致梯度状态异常，间接引起显存问题。标准写法是每个训练脚本只创建一个 scaler，在optimizer.step()的位置用scaler.step()替代。
第四个坑是 Jupyter Notebook 里跑训练。Notebook 的 cell 会保留所有输出和变量引用，你跑一次训练，所有中间张量都被 cell 的_变量持有。这种情况建议把训练逻辑写成.py脚本，用命令行跑，别在 Notebook 里长时间训练。
7. 预防胜于治疗：写训练脚本时的显存纪律
与其等泄漏了再查，不如在写代码时就养成几个习惯。第一，所有累加统计量一律用.item()，绝不直接累加张量。第二，所有需要保存的中间特征一律.detach()，并且限制保存数量。第三，验证和推理一律加@torch.no_grad()。第四，训练循环里定期打印 allocated，发现异常立即停下排查，别等到 OOM 才回头找。
另外，torch.cuda.memory_summary()能打印一份详细的显存报告，包括各尺寸块的分配情况、碎片率等。我在每个项目收尾时都会跑一次，确认没有异常的大块占用。这个报告对判断是碎片问题还是真泄漏很有帮助。
最后分享一个实用技巧：如果你怀疑是某个第三方库导致的泄漏，可以用torch.cuda.memory._record_memory_history()配合调用栈过滤，只看那个库相关的分配。PyTorch 的 snapshot 里每个分配点都有完整的 Python 调用栈，直接搜库名就能定位。这个方法帮我揪出过好几个库内部的缓存 bug，比读源码快得多。
