3. PyTorch显存管理

做服务压测的时候，显存就是你的命根子。我见过太多人模型能跑，一上压测就OOM。说白了，你根本不知道PyTorch在背后怎么折腾显存的。今天咱们就把这事掰扯清楚。

3.1 torch.cuda机制：你以为是自动的？

很多人觉得，我调个.cuda()就完事了。其实没那么简单。PyTorch的CUDA机制，本质上是维护了一个缓存分配器。

我在项目中遇到过一个问题：压测跑着跑着，显存占用越来越高，但明明没有内存泄漏。后来发现是PyTorch的缓存机制在作祟。你可以用下面这行代码看看真实情况：

import torch

# 查看当前分配器状态
print(torch.cuda.memory_summary())

# 清空缓存（但不释放给系统）
torch.cuda.empty_cache()

# 查看实际占用
allocated = torch.cuda.memory_allocated()
cached = torch.cuda.memory_reserved()
print(f"已分配: {allocated / 1024**2:.2f} MB")
print(f"已缓存: {cached / 1024**2:.2f} MB")

你想想看，memory_allocated是真正在用的，memory_reserved是PyTorch占着茅坑不拉屎的。两者差距越大，说明浪费越严重。

3.2 张量存储与释放：别让张量赖着不走

张量的生命周期，说白了就是引用计数。没人引用它了，PyTorch就会回收。但问题在于——很多人不知道什么时候该释放。

我举个例子：

def inference(batch):
    # 坏习惯：创建了中间张量但不释放
    x = batch.cuda()
    hidden = model.layer1(x)
    output = model.layer2(hidden)
    # hidden还活着，占着显存
    return output

# 好习惯：显式释放
def inference_optimized(batch):
    x = batch.cuda()
    output = model.layer2(model.layer1(x))
    # 中间结果随用随丢
    return output

为什么会这样？因为hidden这个变量还在作用域里，PyTorch不敢回收。你把它写成链式调用，中间结果就没了引用，自然被回收。

3.3 梯度累积与显存：反向传播的隐形杀手

做压测的时候，很多人为了模拟真实训练场景，会开梯度累积。嗯，这里要注意——梯度累积会吃掉大量显存。

原因很简单：PyTorch默认会保留所有中间激活值，用于反向传播。你累积的步数越多，保留的激活值就越多。

累积步数	显存占用（相对值）	建议场景
1步	1x	显存充足时
4步	约3.5x	需要小心
8步	约7x	容易OOM

我建议的做法是：

# 显存友好的梯度累积
accumulation_steps = 4
optimizer.zero_grad()

for i, batch in enumerate(dataloader):
    output = model(batch)
    loss = loss_fn(output, target)
    
    # 关键：除以累积步数
    loss = loss / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
        
        # 清空计算图
        torch.cuda.empty_cache()

3.4 torch.no_grad()优化：推理时的必备武器

做服务压测，99%的情况是推理，不是训练。但很多人写代码时，推理和训练混着来。这就出问题了。

torch.no_grad()的作用，说白了就是告诉PyTorch：别给我算梯度，别给我保留中间变量。这样显存占用能降30%-50%。

# 错误示范：推理时还开着梯度
def bad_inference(x):
    with torch.enable_grad():  # 默认就是开启
        output = model(x)
        # 中间激活值全保留着
    return output

# 正确示范：用no_grad包裹
def good_inference(x):
    with torch.no_grad():
        output = model(x)
        # 不保留中间激活值，显存立刻释放
    return output

你想想看，一个ResNet-50，推理时中间激活值能占200MB以上。你用no_grad，这200MB就省下来了。压测时并发一高，差距就出来了。

3.5 实战总结：显存优化的三板斧

说了这么多，我总结一下压测时最实用的三个技巧：

1. 监控先行：用torch.cuda.memory_summary()实时看显存，别等OOM了才慌。
2. 及时释放：中间变量用完就删，链式调用比临时变量好。
3. 推理用no_grad：这是性价比最高的优化，一行代码省30%显存。

最后送大家一张图，把今天讲的逻辑串起来：

这张图把今天讲的四个知识点串起来了。从张量创建开始，经过CUDA分配器，到计算图构建，然后根据是否训练走不同分支，最后统一到显存释放策略。你压测时照着这个流程排查，基本不会出大问题。