汤不热吧在深度学习模型的训练和推理过程中，尤其是在使用PyTorch时，我们经常会遇到一个棘手的问题：明明通过 nvidia-smi 看到显存（GPU Memory）还有剩余，但在尝试分配新的大张量时却报出了 OOM（Out of Memory）错误。这通常是显存碎片化导致的。
本文将深入解析PyTorch的 CachingAllocator 机制，并提供实操代码，教你如何利用 torch.cuda.empty_cache() 来有效缓解这一问题。
1. 为什么会发生显存碎片化？
如果PyTorch每次需要显存都调用CUDA API的 cudaMalloc()，并在释放时调用 cudaFree()，性能会非常慢。为了提高效率，PyTorch引入了一个缓存分配器 (CachingAllocator)。
CachingAllocator 工作原理：
- 当你请求一块显存时，PyTorch首先查看它的内部缓存池。如果找到一个大小合适的已释放块，则直接返回使用（无需调用昂贵的 cudaMalloc）。
- 当你释放显存时（例如张量超出作用域或被 del），显存并不会立刻返回给CUDA驱动，而是被放回 CachingAllocator 的缓存池中。
虽然这极大提高了性能，但随着程序运行，大量的不同大小的张量被分配和释放，缓存池中就会形成许多被小块内存填充的“空隙”。当程序需要一块非常大的连续显存时，即使所有缓存块的总和大于需求量，但因为找不到一块足够大的连续空间，分配就会失败，导致OOM。
2. 深入理解 empty_cache() 的作用
torch.cuda.empty_cache() 是解决碎片化的核心工具。它的作用非常直接：
通知 CachingAllocator 将所有当前处于缓存池中、但未被任何张量引用的空闲内存块，全部释放回给 CUDA 驱动。
请注意，empty_cache() 不会释放当前正在使用的（即有张量引用的）显存。
3. 实操演示：监控缓存和使用 empty_cache()
我们将通过一个简单的PyTorch脚本来观察 CachingAllocator 的行为，以及 empty_cache() 带来的效果。
我们使用 torch.cuda.memory_stats() 来获取详细的显存信息，重点关注 allocated_bytes (实际使用的内存) 和 cached_bytes (缓存池中的内存)。
import torch
import gc
def print_memory_status(message):
allocated = torch.cuda.memory_allocated() / (1024**3)
cached = torch.cuda.memory_reserved() / (1024**3)
print(f"\n--- {message} ---")
print(f"Allocated Memory: {allocated:.2f} GB")
print(f"Cached Memory: {cached:.2f} GB")
if torch.cuda.is_available():
# 1. 初始化和查看初始状态
device = torch.device("cuda")
print_memory_status("初始状态")
# 2. 模拟大量分配和释放，制造缓存和碎片
print("\n--- 步骤 2: 分配并释放大量数据 ---")
# 分配一个大张量
big_tensor = torch.randn(2048, 4096, 4096, device=device) # 约 130 GB * 4 bytes/float = 0.5 GB
# 分配并立即删除许多小张量，造成缓存池内有大量小块
small_tensors = []
for i in range(5000):
t = torch.randn(100, 100, device=device)
small_tensors.append(t)
# 释放大张量和所有小张量
del big_tensor
del small_tensors
gc.collect() # Python垃圾回收确保引用解除
# 查看此时的状态：allocated 应该很低，但 cached 很高
print_memory_status("释放张量后（缓存未清空）")
# 3. 使用 empty_cache() 清理缓存
print("\n--- 步骤 3: 调用 empty_cache() ---")
torch.cuda.empty_cache()
# 4. 再次查看状态
# 此时 cached memory 应该大幅下降，内存被释放回给 CUDA 驱动
print_memory_status("调用 empty_cache() 后")
# 5. 尝试重新分配（如果之前因为碎片化分配失败，现在成功的几率会提高）
# 注意：如果 cached memory 降到接近 0，说明缓存已被清理干净。
else:
print("CUDA 不可用")
预期输出解释：
- 释放张量后（缓存未清空）: 你会发现 Allocated Memory (实际使用) 几乎为零，但 Cached Memory (缓存池) 依然很高。这是因为 PyTorch 留着这些内存块以备下次快速使用。
- 调用 empty_cache() 后: Cached Memory 会大幅度下降，接近甚至等于 Allocated Memory。这意味着 PyTorch 已经把那些空闲的内存块交还给了操作系统/CUDA驱动。
4. 什么时候应该使用 empty_cache()？
虽然 empty_cache() 是解决碎片化的有效工具，但频繁使用会降低性能，因为它重新引入了昂贵的 cudaMalloc 调用。
推荐的使用场景包括：
- 推理阶段切换模型： 在加载新模型之前，旧模型的中间激活值和参数已被清除，此时调用 empty_cache() 可以确保新模型拥有最干净、最连续的显存空间。
- 跨任务边界： 在一个训练 epoch 结束后，或者在执行一个内存密集型的预处理步骤之后。
- 捕获 OOM 错误后： 当程序遇到 OOM 但你确认逻辑上内存足够时，尝试在下一个步骤前调用 empty_cache() 往往能解决问题。
正确理解和使用 torch.cuda.empty_cache()，是AI工程师进行 PyTorch 显存优化和模型部署时必备的技能。