汤不热吧许多AI开发者在使用PyTorch进行训练或推理时，经常会遇到一个困惑：当我使用del删除张量后，或者模型明明只占用了几个GB的显存，但通过nvidia-smi查看时，GPU的显存占用率仍然居高不下。本文将深入解析PyTorch的显存分配机制，并指导您如何正确地监控和管理显存。
为什么 nvidia-smi 看到的显存不等于真实占用？
答案在于 PyTorch 默认使用的是缓存分配器 (Caching Allocator)，而不是传统的直接分配/释放机制。
1. 传统分配 (Malloc/Free) 的缺陷
在传统的内存管理中，每次需要显存时，程序都向CUDA驱动发起请求（cudaMalloc），不需要时则立即释放（cudaFree）。然而，cudaMalloc和cudaFree操作都是相对昂贵的，它们涉及内核调用和同步，会显著降低模型训练和推理的性能。
2. PyTorch 缓存分配器的工作原理
为了提高性能，PyTorch的分配器采用了以下策略：
- Reserved Memory (预留显存): 当程序首次需要显存或需要更多显存时，PyTorch会向CUDA驱动请求一大块内存块（例如256MB）。这部分内存被称为“Reserved”（预留），它是nvidia-smi看到的总占用。
- Allocated Memory (已分配显存): PyTorch从Reserved的内存块中切割出小块，分配给具体的张量（Tensor）。这部分是实际被张量占用的内存。
- Caching (缓存): 当张量被del删除或超出作用域时，PyTorch并不会将这部分内存归还给CUDA驱动（即不执行cudaFree）。而是将其标记为“空闲”，放入一个内部的缓存池 (Cache Pool)中，以便下次需要相同大小或更小块内存时直接快速重用。这是零开销的。
因此，nvidia-smi报告的是 Reserved Memory，而您模型实际使用的、或当前活动张量占用的，是 Allocated Memory。
实践：如何观察缓存和预留显存
我们可以通过 PyTorch 提供的 torch.cuda.memory_stats() 或 torch.cuda.memory_summary() 函数来获取详细的内存报告。
以下代码展示了分配、删除张量后，Reserved 显存保持不变，但 Allocated 显存下降的现象：
import torch
# 确保使用GPU
if not torch.cuda.is_available():
print("CUDA不可用")
exit()
device = torch.device("cuda")
print("--- 1. 初始状态 ---")
# 初始状态，PyTorch可能已经预留了一些内存，也可能没有。
print(torch.cuda.memory_summary(device=device, abbreviated=True))
# --- 2. 首次大块分配 ---
print("\n--- 2. 首次大块分配 (Reserved 显存将增加) ---")
tensor_a = torch.randn(1024 * 1024 * 100, device=device) # 约 400MB
# 此时 nvidia-smi 会看到占用增加，memory_summary会显示 allocated 和 reserved 都增加
print(torch.cuda.memory_summary(device=device, abbreviated=True))
# --- 3. 删除张量 ---
print("\n--- 3. 删除张量 (Allocated 显存减少，Reserved 显存不变) ---")
del tensor_a
# 触发垃圾回收，确保引用被释放
import gc
gc.collect()
# 注意观察此时 Allocated bytes 减少了，但 Reserved bytes 几乎不变 (因为内存进入了缓存池)
print(torch.cuda.memory_summary(device=device, abbreviated=True))
# --- 4. 重新分配 ---
print("\n--- 4. 重新分配 (使用缓存的内存，Reserved 显存不变) ---")
tensor_b = torch.randn(1024 * 1024 * 50, device=device) # 约 200MB
# Allocated bytes 重新增加，但由于使用的是缓存中的空闲块，Reserved bytes 仍然不变。
print(torch.cuda.memory_summary(device=device, abbreviated=True))
关键函数解析
| 函数 | 描述 | 用途 |
|---|---|---|
| torch.cuda.memory_stats() | 返回详细的内存统计字典。 | 程序化监控 |
| torch.cuda.memory_summary() | 返回格式化的人类可读的内存报告。 | 调试和分析 |
| torch.cuda.empty_cache() | 清空缓存池中的所有空闲内存，将其归还给CUDA驱动。 | 释放 Reserved 内存 |
什么时候需要手动清空缓存？
大多数情况下，您不需要手动调用torch.cuda.empty_cache()，因为缓存分配器是为了性能而存在的。
但是，在以下场景中，调用该函数是必要的：
- 进程切换或长时间空闲： 当您完成一个大模型的训练，需要启动另一个完全不同的任务时，可以调用它来释放缓存，确保新任务有最大的可用显存空间。
- OOM (Out of Memory) 错误诊断： 有时模型在运行时遇到OOM，您可能需要在一个更小的批次或更小的模型上测试时，清空之前大模型遗留的缓存。
- 内存碎片化： 虽然PyTorch的分配器会尽量避免碎片化，但在极端动态分配和释放的情况下，如果缓存池中没有连续的足够大的空闲块，即使 Reserved 总量够用，也可能导致OOM。清空缓存可以解决这类问题。
# 示例：释放所有缓存的显存
print("\n--- 5. 清空缓存并释放 Reserved 显存 ---")
torch.cuda.empty_cache()
# 此时 nvidia-smi 看到的占用会大幅下降。
print(torch.cuda.memory_summary(device=device, abbreviated=True))
总结
PyTorch 的缓存分配器是一种性能优化手段，它导致了 nvidia-smi 报告的 Reserved 显存与实际 Allocated 显存不一致。了解这个机制，并使用 torch.cuda.memory_summary() 来监控内部的 Allocated 和 Reserved 状态，是高效管理 GPU 资源的关键。