企业官网建设流程全解析

如果跑过 PyTorch 的大规模训练，你肯定见过这种场景：任务刚起来的时候显存占用很正常，比如 ResNet-50 开 256 的 batch 也就占 4~5GB，但跑着跑着，nvidia-smi里的显存占用像爬坡一样一点点往上涨。从 7GB 涨到 9GB，再涨到 12GB，直到 OOM，甚至影响同机其他卡上的任务。更难受的是，任务本身不报错，loss 还在正常下降，除了显存越占越多，一切仿佛都很正常。这个现象，十有八九就是 PyTorch 训练里的显存泄漏（GPU 内存泄漏）。这篇文章我会结合这些年调训练任务的真实经验，把显存泄漏的成因、定位思路和修复办法一次性说清楚，手里正跑着训练的你绝对用得上。

很多人都把“显存越跑越高”直接等同于“代码写错了”，其实不完全对。PyTorch 有自己的显存缓存机制，nvidia-smi看到的总占用并不等于你实际分配出去的张量占用，有时候只是缓存策略在捣乱。但反过来，要是真泄漏了，你还用“清缓存”的方法去救，那就是南辕北辙。所以第一步必须先把两者分开。

1. 你是真泄漏，还是只是显存本来就不够？

1.1 显存残留和显存泄漏的边界

PyTorch 的 GPU 显存管理走的是 caching allocator 思路。你的代码里tensor.cuda()或tensor.to(device)申请显存，PyTorch 并不会一步到位地直接找 NVIDIA 驱动要，而是先从驱动申请一大块，放进自己的缓存池里，再按需切小块分给你。当你del tensor之后，这块显存只是标记为“空闲”回到缓存池，不一定是真的归还给驱动。所以nvidia-smi看到的总占用会一直保持在一个较高水平，这是正常的。

正常训练时，torch.cuda.memory_allocated()（实际在用的张量内存）会在一个区间内波动，而torch.cuda.memory_reserved()（缓存池总大小）则相对稳定，偶尔因为触发更大的分配而跳一级。真泄漏和这种正常现象的区别在于：泄漏时memory_allocated()会每个 iteration 结束后都比上一个 iteration 高一点，一路累积，不回头。如果你只是memory_reserved()偶尔涨一级，或者nvidia-smi占用高但memory_allocated()稳定，那很可能不是“泄漏”，只是 PyTorch 在帮你缓存显存。

1.2 三分钟快速自检脚本

判断方法其实很简单。在训练循环里，每跑一个 step 打印一次两个指标：

import torch def check_mem(step): torch.cuda.synchronize() allocated = torch.cuda.memory_allocated() reserved = torch.cuda.memory_reserved() print(f"step {step:4d} | alloc {allocated / 1e6:8.1f} MB | reserved {reserved / 1e6:8.1f} MB")

让程序跑上 50~200 步，然后看alloc这一列的变化趋势。如果你发现它每步都多出固定的一截，比如每步多 1.2MB，几十步后已经涨了上百MB，那基本可以判定为真正的显存泄漏。如果alloc始终在一个窄区间里震荡，那就不是泄漏，最多是缓存池增长或有偶发的临时大张量。这时候你可以在关键节点试试torch.cuda.empty_cache()，它会清空缓存池里的空闲块，reserved会明显回落，但allocated不会变——因为它只能清没用的部分。

注意：torch.cuda.max_memory_allocated()是历史峰值，永远只增不减，不能用来判断泄漏。要判断泄漏，请用当前值memory_allocated()。

2. 泄漏的源头，基本逃不过这几种情况

我排查过的显存泄漏问题里，九成都不是玄学，而是下面这几个常见的代码习惯在作妖。每个都看起来不起眼，一旦触发，显存就像断了闸的水龙头。

2.1 训练循环里把 loss 一直当作标量存着

这是最经典、也最容易踩的坑。很多人想画 loss 曲线，于是写了个列表：

loss_history = [] for step, (data, target) in enumerate(train_loader): ... loss = criterion(output, target) loss_history.append(loss) # 错误写法

loss看着像个标量，但它是一个 0 维 Tensor，背后挂着整条反向计算图的引用。你把它塞进列表，计算图就一直没法释放。下一次迭代又生成一张新图，又 append 一次，显存自然线性增长。

正确写法自然是用loss.item()：

loss_history.append(loss.item())

loss.item()返回一个 Python 浮点数，和计算图完全没关系。同理，用 tensorboard 记录 scalar 时也要注意：writer.add_scalar("train/loss", loss, step)在一些版本里会自动取 item，但为了保险，我还是会写成float(loss)。

这个坑在微调 LoRA、跑 YOLOv8，以及自己做 RL 训练记录 reward 时特别常见。凡是“每个 step 都要存一个标量”的地方，检查一遍是否存的是 tensor。

2.2 为了一次 backward 多用几次，就把 retain_graph=True 挂满全场

另一个高频泄漏点是loss.backward(retain_graph=True)。有些代码在梯度惩罚、双阶段反向或者需要计算额外梯度时会这么写：

loss.backward(retain_graph=True) some_other_loss.backward() optimizer.step()

retain_graph=True是告诉 PyTorch 保留反向图，以便后续再次 backward。这个选项本身没毛病，但如果你用的是同一个 loss 对象，完事之后又没有把引用清掉，下一轮迭代又挺着一张旧图，两张图叠一起，显存不涨才怪。

我见过很经典的写法是：训练 GAN 时，判别器和生成器共用一个loss，为了分别 backward 两次，图被保留得特别久。定位方法后面会说，但这里先记住一个原则：能用一次 backward 解决的，绝不要用两次；必须用retain_graph=True时，用完之后立刻del loss，或者把不需要的中间变量置为None，帮助引用计数归零。

2.3 自定义模型里偷懒的 self.xxx = tensor

写自定义 Module 的时候，有些人喜欢在forward里顺便把中间特征保存到实例属性上，方便后面做可视化和调试：

class MyNet(nn.Module): def forward(self, x): ... self.mid_feat = x # 顺手保存一下，方便调试 return x

这个习惯非常危险。self.mid_feat是一个实例属性，它会持有这个 Tensor 的引用，导致计算图释放时，这个激活值仍然被强引用着。每次前向都覆盖一次self.mid_feat，按理说旧的会被替换、释放，但如果mid_feat本身还参与了别的地方的反向，或者你存的是一个 list/字典，那旧引用就一直活着。尤其是如果你还做了self.outputs.append(x)这种操作，那就不是泄漏，而是自爆。

如果确实要保存中间结果，请至少做一步self.mid_feat = x.detach()，并且想清楚它到底占不占 GPU。很多“低显存运行模型”的项目，为了榨干显存，第一步要做的就是清理这类多余的中间保存。

2.4 验证阶段没关梯度，激活值堆成了山

训练时 forward 因为需要反向所以会保留激活值，这正常。但验证/测试的时候，如果你忘了用torch.no_grad()包住，模型依然会以为后续要 backward，于是把每个 batch 的激活值全部缓存下来。一次验证跑完几十个 batch，显存峰值暴涨，看起来就像泄漏。

这个问题在分割、检测这类大输入尺寸模型上尤其致命。跑 mmrotate 训练 DOTA 数据集也好，mmsegmentation 训 Cityscapes 也好，验证阶段一张大图上激活值体积很可观，再叠加“没关梯度”的缓存，显存占用直接翻倍。

正确做法：

with torch.inference_mode(): for data, target in valid_loader: output = model(data) ...

torch.inference_mode()是torch.no_grad()的升级版，更快也更省内存，PyTorch 1.9+ 都能用。如果你还在用老版本，用torch.no_grad()即可。

2.5 钩子和中间结果监听，别遗忘

PyTorch 的register_forward_hook和register_backward_hook是调试和提特征的好工具，但也很容易造成显存泄漏。最常见的是在 hook 里保存输出：

def forward_hook(module, input, output): self.features.append(output) # output 是 GPU tensor

前面说过了，保存 GPU tensor 会引用激活值；而且 hook 注册之后不会自己消失，除非你显式调用handle.remove()。如果这个 hook 是在训练中途注册的，忘记移除，它会在之后每一个 batch 都执行一遍，把输出一股脑全塞进 list。

建议：临时调试用 hook，打印完就handle.remove()；确需保存特征，用output.detach().cpu().numpy()转成 CPU 数据再存，GPU 上不留任何东西。

3. 定位显存泄漏，我用的刀法

知道常见原因之后，最难的是“到底我代码里是哪一个”。毕竟模型结构大、DataLoader、loss、优化器、日志全都纠缠在一起。下面三招是从易到难的排查路径，我自己的经验是：先自检，再二分注释，最后请出 profiler。大部分问题到前两步就能破案。

3.1 搭一个显存观测脚手架

先给训练脚本加一个“血压计”。就在循环里每 N 步打印一次当前 allocated 和 reserved，顺便计算差分：

prev_alloc = None for step, (data, target) in enumerate(train_loader): ... cur = torch.cuda.memory_allocated() if prev_alloc is not None: diff = cur - prev_alloc print(f"step {step}: alloc={cur/1e6:.1f}MB diff={diff/1e6:.2f}MB") prev_alloc = cur

这里有几个关键判断技巧：

• 前 5~10 步的 diff 通常会比较大，因为 CUDA 在热身、cuDNN 在自动调优，忽略即可。
• 如果后 50 步的 diff 基本是正的，而且数值稳定（比如每步 +0.8MB），说明每步都在新增分配。
• 如果 diff 一会儿正一会儿负，但alloc在一个区间里波动，说明只是临时大张量的分配和释放，不是泄漏。

最好把输出重定向到日志文件，让训练继续跑，你过几分钟再回来看趋势。如果alloc曲线一路向上，恭喜，真泄漏坐实了。

3.2 二分注释：锁定嫌疑范围

这是最糙但最有效的方法，灵感来自二分查找。把训练循环拆成几大块：数据加载、模型 forward、loss 计算、backward、optimizer.step、日志/指标保存、验证函数。然后逐个“灭掉”某些块，观察 cur-diff 是否消失。

具体操作时，不要直接删代码，而是注释掉或者用if False跳过。比如怀疑日志保存，就把loss_history.append(...)注释掉再跑 50 步，看 diff 是否下降；如果消失，说明就是它。如果没消失，再把 backward 注释掉，看 forward 本身就涨不涨。这样一轮下来，最多半小时就能锁定到具体某一行。

我在实际排查一个“每 50 个 iteration 多 1.2GB”的泄漏问题时，就是用这个办法。最后发现是验证函数里有个 benchmark 统计，把每个 batch 的预测结果pred直接append到了全局列表，而pred是 GPU tensor，又没有 detach。当我把pred改成pred.detach().cpu().numpy()后，显存曲线立刻变成了直线。

二分注释时建议用 Git 开个临时分支，或者把原始脚本复制一份，避免改来改去忘了哪些是真代码。

3.3 用 profiler 直接看内存分配

如果二分注释也定位不到，或者怀疑是模型内部的某个算子藏了猫腻，就轮到torch.profiler上场。它能给出每个操作分配的显存数值：

from torch.profiler import profile, ProfilerActivity with profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True, record_shapes=True ) as prof: for _ in range(3): train_step() print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

重点看表格里self_cuda_memory_usage这一列，数值最大的操作会成为怀疑对象。不过要注意，profiler 统计的是单个 step 内的分配峰值，对“跨 step 累积型泄漏”帮助有限。它更适合用来排查某个 forward 或 loss 里成百上千次临时张量分配，导致临时显存飙高的问题。

另外，有一个老工具pytorch_memlab的MemReporter可以自动报告逐行/逐张量的显存分配，但这库好久没更新了，跟新版 PyTorch 容易有兼容问题。我一般不会把项目赌在它身上，宁可手写上面那几句话。

3.4 别忘了看 CPU 内存

虽然标题是“显存泄漏”，但我也遇到过表面上是显存越占越高，实际根子是 CPU 内存被吃光，导致 DataLoader 的 worker 卡死、数据没法及时搬到 GPU，最终 GPU 显存被阻塞住的 case。如果nvidia-smi显示显存在涨，free -g也显示 CPU 内存在涨，那问题可能不是 PyTorch 张量，而是某个 list 在存 Python 对象，比如把整个 input/target 都append进去了。

CPU 内存增长可以用 Python 自带的tracemalloc来定位：

import tracemalloc tracemalloc.start() # 跑若干步后 snapshot = tracemalloc.take_snapshot() for stat in snapshot.statistics('lineno')[:20]: print(stat)

它会列出当前每个文件/行创建的 Python 对象占用内存大小，配合显存观测，往往能看出是同一个编码习惯同时污染了 CPU 和 GPU。

4. 修复实操与让训练骨架自带抗体

定位到问题之后，修复一般只需要几行代码。但真正值钱的是后面那套“怎么让以后的项目不再犯”。我自己的做法是建立一套防泄漏的训练骨架，每次开新项目都直接往里面套。

4.1 针对五种泄漏源的“点杀式”修复

把第 2 节里五个坑的修复方案集中列一下，直接照着抄即可：

另外有两个小优化：optimizer.zero_grad(set_to_none=True)比默认的zero_grad()更省显存，因为它直接把梯度置为 None，而不是设为全 0 的 Tensor；还有在自定义 Module 的forward结束后，不需要保留的局部大 Tensor 尽量用del主动释放引用。这不是魔法，但积少成多。

4.2 几个被很多人误解的“清显存”手段

排查过程中你肯定会试这些招，我把我的判断写在这里。

• torch.cuda.empty_cache()：它只清缓存池里的空闲块，不会释放“正在被引用”的 Tensor。真泄漏时调它没用，反而会拖慢训练，因为它把缓存清掉后下一次分配又得重新向驱动申请。它真正的用途是诊断：调用后如果reserved明显下降而allocated不动，说明是缓存问题不是泄漏。
• gc.collect()：强制 Python 回收循环引用，对张量本身不会立刻释放显存，但如果泄漏是因为 Python 对象循环引用导致 Tensor 引用计数没归零，它能帮上忙。我不建议在每个 step 都调，而是在每个 epoch 结束、验证前后可以调一次当保险。
• 重启训练进程：这是排查手段，不是修复手段。如果你发现重启后一两天又涨，说明代码确实有泄漏，别用重启糊弄自己。

关于PYTORCH_CUDA_ALLOC_CONF这个环境变量，有朋友看到显存高就去设置max_split_size_mb或expandable_segments:True。它们主要解决的是显存碎片化，让预留的显存变少，但对累积型泄漏没用。我通常只在allocated稳定但reserved不断变大的时候用这个。

4.3 我在新项目里固定使用的防泄漏训练骨架

直接给一个简化又完整的模板，你可以在这个基础上改：

import torch train_loss = [] for step, (data, target) in enumerate(train_loader): data = data.cuda(non_blocking=True) target = target.cuda(non_blocking=True) output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() optimizer.zero_grad(set_to_none=True) # 所有要保留的数字都只保留数值 train_loss.append(loss.item()) # 显存哨兵：每100步看一眼 if step % 100 == 0: alloc_mb = torch.cuda.memory_allocated() / 1e6 print(f"step {step} | loss {loss.item():.4f} | alloc {alloc_mb:.1f}MB", flush=True) # 验证阶段 model.eval() with torch.inference_mode(): for step, (data, target) in enumerate(valid_loader): output = model(data.cuda(non_blocking=True)) # 不要保存 GPU tensor，除非你确定下一步是 .detach().cpu() pred = output.argmax(dim=1).detach().cpu().numpy()

这套骨架里还有一条隐藏规则：除了模型参数、优化器状态和必要的梯度，GPU 上不应该有任何超过一个 iteration 存活周期的张量。任何你想“暂时存一下”的东西，要么立刻.item()，要么.detach().cpu()，要么别存。

另外，我会在新项目开工第一天就跑一个“显存冒烟测试”：只用 100~500 个样本，训练 200 步，把每步的memory_allocated()记录下来，计算第 50 步和第 150 步的差值。如果在休息后（比如第 100 步后）alloc 还在持续稳步上升，就先停下来定位泄漏再继续训练。这个习惯帮我省下了大量在半夜等 OOM 的时间。

最后再分享一个小经验：真遇到了怎么都查不出来的泄漏，不要只在单卡上排查，分布式训练里的torch.distributed通信原语、broadcast、all_reduce如果被反复创建新的 buffer，也会让显存越跑越高。这时候把NCCL_DEBUG=INFO打开，多看一眼日志里是否有重复的初始化动作，往往会有意外收获。训练脚本稳定之后，给监控脚本加一个告警，当allocated连续 N 个 step 没有回落时就通知你，比盯着nvidia-smi靠谱得多。