第2章：PyTorch显存管理：torch.cuda基础、显存分配与释放机制、torch.cuda.memory_summary()使用
显存管理，说白了就是跟GPU的“内存条”打交道。很多同学跑深度学习模型，一上来就OOM（Out of Memory），然后一脸懵。其实，PyTorch给了我们一套非常强大的工具来管理显存，只是很多人不知道而已。
我个人习惯，在写任何训练脚本之前，先把显存管理这块理清楚。你想想看，如果连模型占了多少显存、哪里泄露了都不知道，那调试起来得多痛苦？
2.1 torch.cuda基础：你每天都在用，但可能没注意
先说说最基础的几个函数。这些我敢打赌，你肯定用过，但未必真正理解它们背后的含义。
2.1.1 设备检测与设置
import torch
# 检查是否有可用的GPU
print(torch.cuda.is_available()) # True or False
# 获取当前使用的GPU索引
print(torch.cuda.current_device()) # 比如 0
# 获取GPU数量
print(torch.cuda.device_count()) # 比如 4
# 设置默认设备
torch.cuda.set_device(0) # 指定使用第0块卡嗯，这里要注意一点：set_device 这个操作是全局的。我在项目中遇到过，有人在一个脚本里调用了 set_device(1)，结果后面所有张量都跑到GPU 1上去了，而他自己还以为在用GPU 0。所以，我建议你尽量用 torch.device('cuda:0') 这种显式指定的方式，而不是依赖全局设置。
2.1.2 张量在CPU与GPU间的搬运
# 创建一个CPU张量
cpu_tensor = torch.randn(3, 3)
# 搬运到GPU
gpu_tensor = cpu_tensor.cuda() # 默认到当前设备
# 或者更推荐的方式：
device = torch.device('cuda:0')
gpu_tensor = cpu_tensor.to(device)
# 从GPU搬回CPU
cpu_tensor_again = gpu_tensor.cpu()这里有个坑，我曾经踩过：频繁地在CPU和GPU之间搬运数据，会严重拖慢训练速度。因为PCIe带宽是有限的，你想想看，每次搬运都要经过PCIe总线，这可比GPU内部计算慢多了。所以，尽量让数据待在GPU上，别来回折腾。
2.2 显存分配与释放机制：PyTorch到底在搞什么鬼？
很多同学以为，调用 del tensor 或者 tensor = None 就能立刻释放显存。其实，PyTorch的显存管理比这复杂得多。
2.2.1 缓存分配器（Caching Allocator）
PyTorch为了加速，使用了一个叫“缓存分配器”的机制。说白了，就是PyTorch向GPU申请显存时，不是用多少申请多少，而是会多申请一些，缓存起来。下次再需要显存时，直接从缓存里拿，省去了向驱动申请的开销。
为什么会这样？因为向GPU驱动申请显存是一个很慢的操作。如果每次创建一个小张量都要去申请，那训练速度会慢得让人抓狂。
核心机制：PyTorch的显存分配器会维护一个缓存池。当你释放一个张量时，显存不会立刻还给操作系统，而是回到缓存池中，供后续使用。
2.2.2 显存碎片化问题
这个缓存机制虽然快，但带来了一个新问题：显存碎片化。我在项目中遇到过，明明总显存还有剩余，但就是分配不出一个连续的大张量。这就是碎片化在作祟。
举个例子：
# 模拟碎片化
a = torch.randn(1000, 1000).cuda() # 分配一个大块
b = torch.randn(500, 500).cuda() # 分配一个中块
del a # 释放大块
c = torch.randn(800, 800).cuda() # 想分配一个中块，但可能失败为什么？因为 a 释放后，那块空间虽然空出来了，但可能被 b 隔开了，导致没有足够大的连续空间给 c。这就是碎片化。
我的建议：尽量一次性分配大张量，避免频繁创建和销毁不同大小的张量。如果实在避免不了，可以考虑使用 torch.cuda.empty_cache() 来清理缓存，但这会带来性能开销。
2.2.3 显存释放的几种方式
| 方式 | 说明 | 效果 |
|---|---|---|
del tensor |
删除Python引用 | 显存回到PyTorch缓存池，不还给OS |
tensor = None |
解除引用 | 同上 |
torch.cuda.empty_cache() |
清空缓存池 | 将缓存池中的显存还给OS |
torch.cuda.reset_peak_memory_stats() |
重置峰值统计 | 不影响显存占用，只重置统计信息 |
注意：empty_cache() 是一个阻塞操作，调用后程序会等待显存真正释放完毕。频繁调用会影响性能，建议只在显存紧张时使用。
2.3 torch.cuda.memory_summary()：你的显存侦探
这个函数是我调试显存问题时的最爱。它就像一台X光机，能把GPU显存里的情况照得一清二楚。
2.3.1 基本用法
import torch
# 在训练循环中插入
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters())
for epoch in range(10):
# ... 训练代码 ...
# 打印显存摘要
print(torch.cuda.memory_summary())
# 或者只打印当前占用
print(f"当前显存占用: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"缓存显存: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")2.3.2 读懂memory_summary的输出
这个函数的输出信息量很大，我来帮你拆解一下：
- Allocated memory：当前实际使用的显存（张量、模型参数等）
- Reserved memory：PyTorch向GPU申请的总显存（包括缓存）
- Peak memory：历史峰值占用
- Large blocks：大块显存的分配情况
- Segment sizes：不同大小段的显存块分布
我曾经用这个函数抓到一个显存泄露的bug。当时模型跑着跑着，显存占用一直在涨，但明明没有新增任何张量。一查 memory_summary，发现有个循环里创建了临时张量，但没有及时释放。加上 del 之后，问题就解决了。
2.3.3 实战：定位显存瓶颈
def check_memory():
"""显存检查小工具，我习惯在每个epoch结束时调用"""
allocated = torch.cuda.memory_allocated() / 1024**3
reserved = torch.cuda.memory_reserved() / 1024**3
peak = torch.cuda.max_memory_allocated() / 1024**3
print(f"已分配: {allocated:.2f} GB")
print(f"已保留: {reserved:.2f} GB")
print(f"峰值: {peak:.2f} GB")
# 如果保留显存远大于已分配，说明缓存太多
if reserved > allocated * 1.5:
print("⚠️ 缓存过多，建议调用 empty_cache()")
torch.cuda.empty_cache()避坑指南：我曾经在分布式训练中忘记在每个进程里单独调用 memory_summary，结果只看到主卡的显存情况，忽略了其他卡的异常。记住，多卡训练时每张卡都要单独检查。
2.4 知识体系总览
下面这张图，是我自己总结的显存管理知识体系。你可以把它当作一个快速索引，遇到问题时回来看看。
嗯，到这里，PyTorch显存管理的基础知识就讲完了。记住，显存管理不是玄学，它是有规律可循的。你只要掌握了 torch.cuda 的基础操作，理解了缓存分配器的机制，再学会用 memory_summary 来诊断问题，大部分显存问题都能迎刃而解。
我刚开始学的时候也踩过不少坑，但后来发现，只要养成定期检查显存的习惯，很多问题都能提前发现。别等到OOM了才去救火，那时候就晚了。