PyTorch显存管理机制：torch.cuda API、分配策略与缓存分配器

聊到显存溢出，很多人第一反应就是“模型太大，显存不够”。但我在实际项目中踩过不少坑，发现很多时候并不是显存真的不够，而是我们没用好PyTorch的显存管理机制。说白了，PyTorch有一套自己的“小算盘”，你摸透了它，就能省下不少显存。

这一章，我就带你看看PyTorch到底是怎么管显存的。我会从torch.cuda API讲起，再到分配与释放策略，最后深入缓存分配器的工作原理。嗯，这部分内容有点底层，但搞懂了，你排查显存问题会顺手很多。

torch.cuda API：你的显存“遥控器”

PyTorch提供了一套完整的CUDA API，说白了就是让你能直接跟GPU显存对话。我个人习惯在调试时频繁使用这些API，它们能帮我快速定位问题。

先看几个最常用的：

import torch

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

# 查看显存总量（字节）
print(torch.cuda.get_device_properties(0).total_memory)

# 查看已分配显存（字节）
print(torch.cuda.memory_allocated())

# 查看缓存中的显存（字节）
print(torch.cuda.memory_reserved())

# 清空缓存
torch.cuda.empty_cache()

这里有个细节我要强调一下：memory_allocated() 和 memory_reserved() 的区别。前者是实际被张量占用的显存，后者是PyTorch向GPU申请的总显存（包括缓存）。我在项目中遇到过有人只看memory_allocated，结果发现显存占用不大，但程序还是OOM了——其实就是缓存占了大头。

显存分配与释放策略：PyTorch的“小算盘”

PyTorch的显存分配策略，说白了就是“懒分配，勤缓存”。它不会每次需要显存都去跟GPU驱动要，而是先申请一大块，然后自己内部管理。

具体来说：

• 分配时： 当你创建一个CUDA张量，PyTorch会先检查缓存里有没有足够大的空闲块。有就直接用，没有才向GPU申请新的显存。
• 释放时： 当张量被销毁，显存不会立刻还给GPU驱动，而是被PyTorch缓存起来，留给后续使用。

为什么会这样？因为向GPU驱动申请显存是个开销很大的操作。PyTorch这么做，说白了就是为了性能。我刚开始做训练时，发现显存占用只增不减，还以为有内存泄漏，后来才知道是缓存机制在作祟。

举个例子：

import torch

# 创建一个1GB的张量
x = torch.randn(256, 1024, 1024, device='cuda')
print(f"分配后: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"缓存后: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

# 删除张量
del x
print(f"删除后分配: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"删除后缓存: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

你会发现，删除张量后，memory_allocated 降为0，但 memory_reserved 基本没变。这就是缓存的效果。

缓存分配器工作原理：PyTorch的“内存管家”

PyTorch的缓存分配器，本质上是一个“大管家”。它管理着一大块显存，负责分配和回收。它的核心数据结构是“空闲块列表”和“已分配块列表”。

我画了一张图，帮你理解它的工作流程：

这个分配器有几个关键特性：

• 最佳适配策略： 分配时，它会从空闲块列表中找到大小最合适的块，避免浪费。
• 块合并： 释放时，如果相邻的空闲块可以合并，它会合并成更大的块。
• 预分配： 第一次申请显存时，PyTorch会一次性申请一大块（默认2MB），后续小分配都在这个块里切分。

我记得有一次，我在训练一个超大模型，发现显存碎片化严重。明明总空闲显存还有2GB，但就是分配不出一个1GB的连续块。这就是因为缓存分配器里的空闲块都是零散的，合并策略没来得及生效。

实战：如何利用这些知识排查显存溢出

知道了这些机制，我们就能有针对性地排查问题了。我总结了几条经验：

1. 先看缓存，再看分配： 用 memory_reserved() 和 memory_allocated() 对比，如果差值很大，说明缓存占用了大量显存。
2. 检查是否有未释放的张量： 用 torch.cuda.memory_summary() 查看哪些张量还活着。我曾经发现一个循环里忘了删除中间变量，导致显存持续增长。
3. 合理使用 empty_cache()： 在模型保存或评估前调用一次，释放缓存。但不要在训练循环里频繁调用。
4. 利用 torch.no_grad()： 推理时关闭梯度计算，能减少大量中间变量的显存占用。

举个例子，假设你发现训练时显存不断增长，可以这样排查：

import torch

# 在训练循环中打印显存状态
for epoch in range(10):
    for batch in dataloader:
        # 训练代码...
        
        # 打印显存
        allocated = torch.cuda.memory_allocated() / 1024**3
        reserved = torch.cuda.memory_reserved() / 1024**3
        print(f"Epoch {epoch}, Batch {batch_idx}: allocated={allocated:.2f}GB, reserved={reserved:.2f}GB")
        
        # 如果发现allocated持续增长，说明有张量没释放
        # 如果reserved增长但allocated稳定，说明缓存正常

嗯，这里要注意：如果 allocated 持续增长，那大概率是代码里有张量没被释放。检查一下循环里有没有创建了张量但没赋给变量，或者有没有把张量存在列表里忘了清空。

好了，这一章的内容就到这里。PyTorch的显存管理机制其实不复杂，但理解它能让你的调试效率翻倍。下一章我们会聊更具体的显存优化技巧，到时候见。