3. PyTorch显存管理机制

说到PyTorch的显存管理，我刚开始接触时也踩过不少坑。很多人以为调用了torch.cuda.empty_cache()就能清空显存，结果发现根本没用——嗯，这里面的门道，我们得从PyTorch的分配策略说起。

3.1 PyTorch显存分配策略

PyTorch的显存分配，说白了就是「先借后还」的模式。它不会每次需要显存都去问GPU要，而是提前申请一大块，自己内部管理。

我个人习惯把这种策略叫做「批发零售」模式：

批发：PyTorch一次性向CUDA驱动申请大块显存（比如2GB）
零售：内部再把这些大块切成小块，分配给各个Tensor

为什么要这么设计？因为每次调用cudaMalloc都是有开销的。我在项目中遇到过，如果频繁申请释放小张量，性能会下降30%以上。所以PyTorch干脆搞了个「缓存池」，用完了不急着还，留着下次用。

核心机制：PyTorch的显存分配是「懒惰分配 + 缓存复用」的。你创建Tensor时，它可能还没真正占用显存，只有当你开始计算时，才会触发实际分配。

3.2 缓存分配器（Caching Allocator）

缓存分配器是PyTorch显存管理的核心。它维护了一个「已分配块」和「空闲块」的列表。你想想看，当你的代码执行tensor.cuda()时，背后发生了什么？

# 伪代码示意
def allocate_tensor(size):
    # 1. 先查缓存池有没有合适大小的空闲块
    block = caching_allocator.find_free_block(size)
    if block:
        return block  # 直接复用
    
    # 2. 没有的话，向CUDA申请新块
    new_block = cudaMalloc(size)
    caching_allocator.add_block(new_block)
    return new_block

这里有个关键点：缓存分配器是按「大小类别」来管理块的。比如256字节、512字节、1KB、2KB...每个类别都有独立的空闲链表。这样做的好处是分配速度快，但坏处是——容易产生碎片。

我的经验：如果你发现显存占用忽高忽低，很可能是缓存分配器在「整理房间」。这时候调用torch.cuda.empty_cache()，它会释放所有空闲块，但已分配的Tensor不会受影响。

3.3 显存碎片化问题

显存碎片化，说白了就是「大块显存被切成小块，导致无法分配连续大块内存」。我曾经遇到一个场景：明明还有4GB空闲显存，但创建个3GB的Tensor却报OOM。为什么？因为那4GB是分散的，没有连续3GB的空间。

碎片化主要有两种：

内部碎片：分配器为了对齐，多分配了一些空间。比如你要257字节，它可能给你分配512字节的块。
外部碎片：频繁申请释放不同大小的Tensor，导致空闲块不连续。

避坑指南：我曾经在训练动态图模型时，每步都创建不同大小的中间变量，结果显存碎片率高达40%。后来我改用固定大小的batch和序列长度，碎片问题就缓解了很多。

怎么检查碎片化程度？可以用这个命令：

# 查看当前显存使用情况
print(torch.cuda.memory_summary())

# 重点关注这几项：
# - Allocated: 实际使用的显存
# - Cached: 缓存分配器持有的显存（包括空闲块）
# - Fragmentation: 碎片化比例

3.4 torch.cuda系列API详解

这些API我几乎每天都要用。咱们挑几个最实用的说说：

API	作用	我的使用场景
`torch.cuda.memory_allocated()`	返回当前已分配的显存（字节）	监控训练过程中的显存峰值
`torch.cuda.memory_cached()`	返回缓存分配器持有的总显存	判断是否有「隐藏」的显存占用
`torch.cuda.empty_cache()`	释放所有空闲缓存块	在模型切换或评估前清理一次
`torch.cuda.reset_peak_memory_stats()`	重置峰值统计	开始新一轮监控前清零
`torch.cuda.set_per_process_memory_fraction()`	限制进程最大显存使用比例	多进程训练时防止互相抢占

这里有个小技巧：

# 在训练循环中监控显存
def train_step():
    torch.cuda.reset_peak_memory_stats()
    
    # ... 训练代码 ...
    
    peak = torch.cuda.max_memory_allocated() / 1024**3
    current = torch.cuda.memory_allocated() / 1024**3
    print(f"峰值: {peak:.2f}GB, 当前: {current:.2f}GB")

个人习惯：我一般会在每个epoch结束后调用一次empty_cache()，但不会在训练过程中频繁调用——因为清理缓存本身也有开销，而且会破坏分配器的「热缓存」状态。

最后说一个容易忽略的API：torch.cuda.memory_stats()。它会返回一个字典，包含非常详细的分配信息。我在排查显存泄漏时，经常用它来对比不同时间点的分配情况。

stats = torch.cuda.memory_stats()
# 重点关注
print(stats['num_alloc_retries'])  # 分配重试次数，越高说明碎片越严重
print(stats['segment.all.current'])  # 当前段数量

嗯，PyTorch的显存管理机制，说白了就是「用空间换时间」——多占点缓存，换来更快的分配速度。理解了这个核心思想，很多显存问题就能想通了。

这张图把整个流程串起来了。用户代码请求显存 → 缓存分配器先查缓存池 → 有就直接复用，没有就找CUDA驱动申请 → 申请到的块加入缓存池管理。碎片化问题就出在「频繁分配释放不同大小块」这个环节上。

总结一句话：PyTorch的显存管理，核心是「缓存复用」四个字。理解了这个，你就能明白为什么empty_cache()不能真正释放显存，为什么显存碎片化会导致OOM，以及怎么用API来监控和优化显存使用。

想要更多资料，联系可可

📚 课程目录

3. PyTorch显存管理机制

3.1 PyTorch显存分配策略

3.2 缓存分配器（Caching Allocator）

3.3 显存碎片化问题

3.4 torch.cuda系列API详解