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} fallocated{allocated:.1f}MB freserved{reserved:.1f}MB fmax{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 降下来。解决办法是尽量固定输入 shapepadding 到固定长度、避免在训练循环里创建临时大张量。3. 计算图没释放最常见的泄漏源头3.1 计算图持有机制回顾PyTorch 的动态图在loss.backward()之后理论上计算图会被释放。但有几个常见操作会让图一直被持有把loss存进了列表或字典比如losses.append(loss)用于后续统计。这个loss是带 grad_fn 的张量它持有整条计算图的引用。在backward()之前做了loss.item()之外的张量操作比如total_loss loss而不是total_loss loss.item()。用了retain_graphTrue但没在合适时机释放。我见过最典型的写法是这样的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, topk5): 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 losstotal 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 的闭包和全局变量是显存泄漏的隐形杀手。比如你在函数里定义了一个 hookhook 里引用了某个中间张量而这个 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_memoryTrue配合non_blockingTrue能加速传输但要注意 pin memory 本身也占内存。4.3 优化器状态和 EMAAdam 优化器会为每个参数维护 exp_avg 和 exp_avg_sq 两个状态显存占用是参数量的两倍。如果你在训练中途重建了优化器但没释放旧的或者用了 EMA指数移动平均却忘了 detach都会导致显存增长。EMA 的正确写法class EMA: def __init__(self, model, decay0.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_entries100000) # 跑一段训练 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(enabledNone)导出的 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_workersFalse中途重建模型后显存翻倍旧模型未释放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比读源码快得多。