3. PyTorch显存管理：torch.cuda API详解、显存分配策略、缓存机制与清空

显存管理，说白了就是跟GPU要内存、用内存、还内存的过程。很多同学跑模型时遇到OOM（Out of Memory），第一反应就是“模型太大了”，其实很多时候是显存管理没做好。我刚开始调大模型时也吃过这个亏，明明显存够用，就是莫名其妙爆了。

这一节，咱们把PyTorch的显存管理机制彻底讲透。你掌握了这些，以后排查显存问题会顺手很多。

3.1 torch.cuda API：你的显存工具箱

PyTorch提供了一套完整的torch.cuda接口，用来监控和管理显存。我个人习惯把这些API分成三类：查询类、控制类、调试类。

3.1.1 查询类API

先看最常用的几个：

import torch

# 查看当前设备
print(torch.cuda.current_device())

# 查看设备名称
print(torch.cuda.get_device_name(0))

# 查看显存总量（字节）
total_mem = torch.cuda.get_device_properties(0).total_memory
print(f"总显存: {total_mem / 1024**3:.2f} GB")

# 查看已分配显存（字节）
allocated = torch.cuda.memory_allocated(0)
print(f"已分配: {allocated / 1024**3:.2f} GB")

# 查看缓存显存（字节）
cached = torch.cuda.memory_reserved(0)
print(f"缓存: {cached / 1024**3:.2f} GB")

# 查看当前峰值显存
peak = torch.cuda.max_memory_allocated(0)
print(f"峰值: {peak / 1024**3:.2f} GB")

这里有个细节要注意：memory_allocated是实际被张量占用的显存，而memory_reserved是PyTorch向GPU申请的缓存池大小。两者之差就是“空闲但没还给GPU”的显存。

3.1.2 控制类API

这类API用来手动管理显存的分配和释放：

# 清空缓存（释放缓存池中的空闲块）
torch.cuda.empty_cache()

# 重置峰值统计
torch.cuda.reset_peak_memory_stats()

# 设置当前设备
torch.cuda.set_device(0)

empty_cache()这个函数，很多人以为它能释放所有显存。其实不是的——它只释放缓存池中那些“空闲”的块，正在被张量占用的显存它动不了。我曾经在调试一个模型时，发现调用empty_cache()后显存只降了一点点，当时还以为是API坏了，后来才明白原理。

3.1.3 调试类API

排查显存问题时，这些API能帮你定位到具体是哪行代码在吃显存：

# 开启内存快照（用于调试）
torch.cuda.memory._record_memory_history()

# 查看当前所有张量的分配信息
print(torch.cuda.memory_summary())

# 查看某个张量的存储位置
x = torch.randn(1000, 1000, device='cuda')
print(x.device)          # cuda:0
print(x.storage().data_ptr())  # 显存地址

3.2 显存分配策略：PyTorch是怎么“要内存”的

PyTorch的显存分配策略，说白了就是一个“缓存池”机制。它不会每次创建张量都去跟GPU驱动要内存，那样太慢了。而是提前申请一大块，然后自己管理。

这个策略的核心思想是：减少与驱动层的交互次数。你想想看，每次调用cudaMalloc都是系统调用，开销很大。PyTorch的做法是：

1. 预分配：第一次创建张量时，PyTorch会向GPU申请一块较大的内存（默认是2MB的倍数）
2. 缓存复用：释放张量后，内存不还给GPU，而是留在缓存池里
3. 按需分割：新张量需要内存时，从缓存池中切一块出来
4. 合并碎片：当缓存池中有多个相邻空闲块时，PyTorch会尝试合并

我用一个简单的例子来说明：

import torch

# 第一次分配：PyTorch会向GPU申请一大块
a = torch.randn(1000, 1000, device='cuda')
print(f"分配后: allocated={torch.cuda.memory_allocated()/1024**2:.2f}MB, "
      f"cached={torch.cuda.memory_reserved()/1024**2:.2f}MB")

# 释放a：内存回到缓存池，但没还给GPU
del a
print(f"释放后: allocated={torch.cuda.memory_allocated()/1024**2:.2f}MB, "
      f"cached={torch.cuda.memory_reserved()/1024**2:.2f}MB")

# 再分配b：直接从缓存池拿，不需要跟GPU交互
b = torch.randn(500, 500, device='cuda')
print(f"再分配: allocated={torch.cuda.memory_allocated()/1024**2:.2f}MB, "
      f"cached={torch.cuda.memory_reserved()/1024**2:.2f}MB")

运行结果你会看到：释放后allocated降了，但cached没变。这就是缓存机制在起作用。

3.3 缓存机制：为什么显存“只增不减”

很多同学发现，自己的程序跑完后，用nvidia-smi一看，显存还是被占着。这不是泄漏，而是PyTorch的缓存机制在“作祟”。

PyTorch的缓存分配器（Caching Allocator）维护了一个空闲块列表。当张量被释放时，内存块会回到这个列表里。下次分配时，分配器会先在这个列表里找大小合适的块，找不到才去申请新的。

这个机制的好处很明显：

• 分配速度快（避免系统调用）
• 减少碎片（通过合并相邻空闲块）

但坏处也很直接：

• 显存占用看起来比实际需要的大
• 如果缓存池太大，其他进程可能拿不到显存

我曾经在一个共享GPU的服务器上跑实验，我的程序结束后显存还占着2GB没释放，同事跑模型时直接OOM了。后来我养成了习惯，在程序末尾显式调用torch.cuda.empty_cache()。

3.4 清空显存：什么时候该用，什么时候不该用

torch.cuda.empty_cache()这个函数，用得好是神器，用不好是累赘。我总结了几条经验：

3.4.1 该用的时候

• 程序结束前：释放缓存，方便其他进程使用
• 切换模型时：比如你在做模型对比实验，加载新模型前清一下缓存
• 显存碎片严重时：如果频繁分配释放不同大小的张量，缓存池可能碎片化，清空后重新分配会更高效

3.4.2 不该用的时候

• 训练循环中频繁调用：每次清空缓存都会导致下次分配时重新申请，反而降低性能
• 期望它释放所有显存：它只释放空闲缓存，正在使用的张量它动不了
• 作为OOM的万能解药：如果显存真的不够用，清空缓存也救不了你

# 错误示范：在训练循环里频繁清空
for epoch in range(100):
    for batch in dataloader:
        # ... 训练代码 ...
        torch.cuda.empty_cache()  # 别这么干！性能会暴跌

# 正确做法：只在需要的时候清
# 比如加载新模型前
model_a = load_model('model_a.pth')
# ... 用model_a做推理 ...
del model_a
torch.cuda.empty_cache()  # 清掉model_a的缓存

model_b = load_model('model_b.pth')

3.5 知识体系总览

下面这张图总结了PyTorch显存管理的核心逻辑，我画出来方便你理解：

这张图把整个流程串起来了。你从用户创建张量开始，顺着箭头走一遍，就能理解PyTorch是怎么管理显存的。核心就一句话：能复用就不新申请，能缓存就不释放。

3.6 实战建议

最后，我把自己多年积累的几条显存管理经验分享给你：

1. 监控先行：在训练脚本里加上显存监控，每轮打印allocated和cached，这样能第一时间发现问题
2. 及时释放：不再需要的张量用del删除，尤其是中间变量和梯度
3. 慎用empty_cache：只在程序切换或显存碎片严重时调用，不要在热循环里用
4. 理解缓存：看到nvidia-smi显存占用高别慌，先查memory_allocated，那才是实际使用的量
5. 善用上下文管理器：比如torch.no_grad()可以省掉推理时的梯度显存

嗯，这一节的内容就到这里。显存管理是性能优化的基础，你把这些API和策略吃透了，后面排查带宽瓶颈时会轻松很多。