4. 显存优化基础：张量生命周期管理、显存碎片化问题、PyTorch显存分配器机制
显存优化，说白了就是搞清楚你的数据在GPU里是怎么待着的，又是怎么被清理出去的。很多同学一上来就盯着batch size调来调去，结果发现显存还是爆了。其实啊，很多时候问题出在张量的生命周期管理上，而不是你显存真的不够用。
我个人习惯把显存优化分成三个层次来看：第一层是搞清楚张量什么时候生、什么时候死；第二层是理解显存碎片化这个隐形杀手；第三层才是摸透PyTorch的显存分配器到底在背后干了什么。今天咱们就把这三件事儿聊透。
4.1 张量生命周期管理
每个张量从创建到销毁，都经历一个完整的生命周期。我见过太多人写代码时，张量创建了一堆，但用完不释放，结果显存越占越多。
核心要点：张量的生命周期 = 创建 → 使用 → 引用计数归零 → 内存回收
在PyTorch里，张量的生命周期受Python引用计数机制控制。当一个张量不再被任何变量引用时，它就会被标记为可回收。但这里有个坑——引用循环。
我曾经在调试一个训练脚本时，发现显存占用只增不减。排查了半天，原来是某个自定义Dataset里，张量被列表引用着，训练完也没清空。你想想看，每个epoch结束，列表里堆满了历史张量，显存不爆才怪。
避坑指南：
- 及时del：用完的大张量，手动del一下，再配合torch.cuda.empty_cache()
- 避免引用循环：检查自定义类里是否有互相引用的张量
- 使用with torch.no_grad()：推理时关闭梯度计算，减少中间张量
# 不好的写法
def train():
for batch in dataloader:
x = batch.cuda()
y = model(x)
loss = criterion(y, target)
# x, y, loss 都还活着，直到下一轮迭代
# 好的写法
def train():
for batch in dataloader:
x = batch.cuda()
y = model(x)
loss = criterion(y, target)
loss.backward()
optimizer.step()
del x, y, loss # 显式释放
torch.cuda.empty_cache() # 清理缓存
小技巧：我习惯在训练循环里加一个显存监控，每100步打印一次torch.cuda.memory_allocated()和torch.cuda.memory_reserved()。这样能第一时间发现显存泄漏。
4.2 显存碎片化问题
显存碎片化，这是个容易被忽视但杀伤力极大的问题。你可能会遇到这种情况：明明显存还有剩余，但一创建新张量就报OOM。嗯，这就是碎片化在搞鬼。
为什么会这样？说白了，就是频繁创建和销毁不同大小的张量，导致显存被切割成很多小块。每个小块单独看都不大，但合起来能用的连续空间却很少。
我记得有一次跑一个多模态模型，输入图片尺寸不固定，每次forward都会创建不同大小的中间张量。跑了几个小时后，显存占用显示还有2GB空闲，但就是分配不出一个512MB的张量。这就是典型的碎片化。
碎片化的成因：
- 动态形状：每次forward张量尺寸不一样
- 频繁创建临时张量：比如在循环里反复创建小张量
- 混合精度训练：fp16和fp32张量交替创建，大小不一
注意：碎片化问题在长时间训练中会越来越严重。我建议每训练一定步数后，手动调用torch.cuda.empty_cache()来整理碎片。虽然这会影响一点性能，但能避免中途OOM。
解决碎片化的几个实用方法：
- 固定输入尺寸：尽量让所有输入张量尺寸一致
- 预分配大块内存：用torch.cuda.caching_allocator_alloc()预分配
- 使用内存池：自己维护一个张量池，重复利用
- 调整分配器配置：设置PYTORCH_CUDA_ALLOC_CONF环境变量
# 设置环境变量来优化分配器
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
# 或者使用更激进的配置
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
4.3 PyTorch显存分配器机制
PyTorch的显存分配器，说白了就是一个缓存池。它不会每次创建张量都去操作系统申请内存，而是先预申请一大块，然后自己管理。
这个机制的好处是快——避免了频繁的系统调用。但坏处是——你看到的显存占用可能比实际使用的多。因为分配器预申请的内存，即使张量释放了，也不会立刻还给系统。
我刚开始做优化时，看到nvidia-smi显示显存占用80%，但实际模型只用了60%，剩下的20%就是分配器缓存。当时我还以为是内存泄漏，查了半天才发现是分配器在搞鬼。
分配器的核心机制：
- 缓存机制：释放的张量内存不还给OS，而是放入缓存池
- 分块管理：将显存分成不同大小的块，按需分配
- 合并策略：相邻的空闲块会被合并成大块
- 分割策略：大块可以被分割成小块使用
| 分配器参数 | 说明 | 推荐值 |
|---|---|---|
| max_split_size_mb | 控制最大分割块大小 | 128-512 |
| expandable_segments | 启用可扩展段 | True |
| roundup_power2_divisions | 分配对齐方式 | 256 |
关键理解：torch.cuda.memory_allocated()返回的是实际使用的显存，而torch.cuda.memory_reserved()返回的是分配器预申请的总显存。两者之差就是缓存池的大小。
这里我画了一张图，帮你理解整个显存管理的流程：
你看这个流程，张量创建时先问分配器有没有缓存块，有就直接用，没有才去操作系统申请。释放时也不还给OS，而是放回缓存池。这个设计在大多数场景下是高效的，但如果你不注意张量生命周期，缓存池里就会堆满各种大小的碎片块。
实用建议：我一般会在训练脚本开头设置PYTORCH_CUDA_ALLOC_CONF，把max_split_size_mb设成128或256。这样能减少大块被分割成小块的概率，从而降低碎片化。另外，如果显存比较紧张，可以试试expandable_segments:True，这个选项能让分配器更灵活地管理内存。
最后说一句，显存优化不是一蹴而就的事。我建议你先用torch.cuda.memory_summary()打印一份完整的显存报告，看看分配器到底在干什么。然后再针对性地调整张量生命周期和分配器配置。这样一步步来，显存利用率能提升不少。