3、显存泄漏的修复策略

显存泄漏，说白了就是内存只进不出。你训练一个模型，显存占用像滚雪球一样越滚越大，直到OOM崩掉。我早期做项目时就被这个问题折磨过——明明模型不大，跑着跑着就爆了。后来才发现，不是模型大，是代码写得不够干净。

这一章，我把自己踩过的坑和修复策略整理出来。每个方法我都亲自验证过，希望能帮你少走弯路。

3.1 with torch.no_grad() 上下文管理

推理阶段，最常用的就是 torch.no_grad()。为什么？因为推理不需要计算梯度。如果你忘了关梯度，计算图会一直累积，显存自然就炸了。

核心原则：推理时一定要用 with torch.no_grad() 包裹。

# 错误示范：没有关闭梯度
for batch in dataloader:
    output = model(batch)  # 计算图保留，显存泄漏

# 正确做法
with torch.no_grad():
    for batch in dataloader:
        output = model(batch)  # 不保留计算图，显存安全

我个人习惯在验证和测试阶段，直接在最外层套上 torch.no_grad()。这样即使内部有复杂的操作，也不会意外保留计算图。

小技巧：如果你用 @torch.no_grad() 装饰器装饰函数，也能达到同样效果。但我更喜欢显式的 with 语句，因为作用域更清晰。

3.2 del 与 torch.cuda.empty_cache() 配合使用

有时候，即使你关了梯度，显存还是涨。为什么？因为Python的垃圾回收不是实时的。你删除了变量，但显存并没有立即释放。

这时候，del 和 torch.cuda.empty_cache() 就派上用场了。

# 手动释放显存
large_tensor = torch.randn(10000, 10000).cuda()
# ... 使用 large_tensor ...
del large_tensor  # 删除引用
torch.cuda.empty_cache()  # 清空缓存

我曾经在做一个多模型推理服务时，发现每次切换模型，显存都会涨一点。排查了半天，发现是旧模型的参数没有被显式删除。加上 del model 和 torch.cuda.empty_cache() 后，显存就稳定了。

注意：empty_cache() 不是万能的。它只释放缓存，不释放正在使用的显存。如果你还有引用指向某个张量，调用 empty_cache() 是没用的。

3.3 梯度清零的最佳实践

训练时，梯度是累积的。如果你不清零，上一轮的梯度会一直占用显存。PyTorch 提供了 optimizer.zero_grad()，但很多人用错了地方。

# 标准训练循环
for batch in dataloader:
    optimizer.zero_grad()  # 每个batch开始前清零
    output = model(batch)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

嗯，这里要注意：zero_grad() 的位置很关键。放在 backward() 之后？不行，因为梯度还没用。放在 step() 之后？也可以，但容易漏掉。我个人习惯放在每个batch的开头，这样逻辑最清晰。

避坑指南：我曾经在分布式训练中，忘了在 loss.backward() 之前同步梯度。结果每个进程的梯度都不一样，模型直接不收敛。后来加上了 torch.distributed.all_reduce 才解决。

3.4 DataLoader 的 num_workers 与 pin_memory 正确配置

DataLoader 的配置直接影响显存。很多人觉得 num_workers 越大越好，其实不然。

参数	作用	显存影响
num_workers	数据加载的进程数	每个worker会复制数据，显存占用翻倍
pin_memory	将数据固定在内存中	加速数据传输，但占用更多内存

你想想看，如果 num_workers=8，每个worker都加载一批数据到显存，那显存占用就是8倍。我建议：

小模型：num_workers 可以设大一点（4-8），因为显存充裕
大模型：num_workers 设小一点（2-4），避免显存爆炸
pin_memory：如果显存紧张，可以关掉。虽然慢一点，但更安全

# 显存紧张时的配置
dataloader = DataLoader(
    dataset,
    batch_size=16,
    num_workers=2,  # 减少worker数
    pin_memory=False,  # 关闭固定内存
    prefetch_factor=2  # 预取因子，控制预加载量
)

个人经验：我一般先设 num_workers=4，然后观察显存变化。如果显存波动大，就降到2。如果CPU负载高，就升到8。没有固定值，得试。

3.5 多进程显存隔离

多进程训练时，显存隔离是个大问题。每个进程都会加载一份模型和数据，如果隔离不好，显存会互相干扰。

PyTorch 的 torch.multiprocessing 提供了 spawn 和 fork 两种启动方式。我强烈建议用 spawn，因为它会创建全新的进程，不会继承父进程的显存状态。

# 推荐：使用 spawn 启动多进程
import torch.multiprocessing as mp

def train(rank, world_size):
    # 每个进程独立初始化
    model = MyModel().cuda(rank)
    # ... 训练逻辑 ...

if __name__ == "__main__":
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

我曾经踩过一个坑：用 fork 启动多进程，结果子进程继承了父进程的CUDA上下文，显存直接翻倍。换成 spawn 后，每个进程的显存就独立了。

注意：使用 spawn 时，所有代码必须放在 if __name__ == "__main__" 里，否则会无限递归创建进程。

知识体系总览

下面这张图，概括了显存泄漏修复的完整策略。你可以把它当作检查清单，每次遇到显存问题，按图索骥就行。

这五个策略，每个都能解决一类显存问题。但实际项目中，往往是多个问题叠加。我建议你遇到显存泄漏时，按这个顺序排查：先看推理有没有关梯度，再看有没有手动释放大张量，然后检查梯度清零的位置，接着优化DataLoader配置，最后确认多进程隔离是否做好。

嗯，这一章的内容就到这里。记住，显存泄漏不是玄学，是代码习惯问题。养成好习惯，显存自然就稳了。

📚 课程目录

3、显存泄漏的修复策略

3.1 with torch.no_grad() 上下文管理

3.2 del 与 torch.cuda.empty_cache() 配合使用

3.3 梯度清零的最佳实践

3.4 DataLoader 的 num_workers 与 pin_memory 正确配置

3.5 多进程显存隔离

知识体系总览