3、PyTorch显存机制：Tensor存储原理，CUDA缓存分配器，显存碎片化问题
聊到显存优化，咱们得先搞清楚一件事：PyTorch到底是怎么管显存的？
很多同学训练到一半就OOM，第一反应就是“batch size调小点”。但说实话，这往往只是治标不治本。我见过太多项目，调小batch size后显存确实降了，但训练速度慢得让人抓狂。真正的问题，可能出在你对PyTorch显存机制的理解上。
今天我就带你把这层窗户纸捅破。咱们从三个层面来拆解：Tensor怎么存、CUDA怎么分配、碎片化怎么来的。
3.1 Tensor存储原理：你以为的“释放”可能只是假象
先问个问题：当你执行 del tensor 时，显存真的被释放了吗？
答案是：不一定。
PyTorch的Tensor存储，底层依赖的是 Storage 对象。每个Tensor背后都有一个Storage，它才是真正占显存的那块“地皮”。Tensor本身只是个“门牌号”，告诉你数据在哪、形状多大。
核心概念：
- Tensor：视图（View），包含元数据（shape、stride、dtype）
- Storage：实际数据块，连续内存区域
- 引用计数：多个Tensor可以共享同一个Storage
举个例子你就明白了：
import torch
# 创建一个Tensor
a = torch.randn(1000, 1000, device='cuda')
print(a.storage().data_ptr()) # 打印Storage地址
# 创建一个视图（切片）
b = a[0:500, :]
print(b.storage().data_ptr()) # 地址相同！共享Storage
# 删除a
del a
# 此时b依然有效，因为Storage的引用计数还没归零
print(b.sum()) # 正常运行
看到了吗？del a 只是减少了引用计数，Storage还在。只有所有引用都消失，Storage才会被标记为“可回收”。
我的经验： 在项目中排查显存泄漏时，我习惯用 torch.cuda.memory_summary() 查看当前所有Tensor的分配情况。这个命令能列出每个Tensor的Storage地址和大小，帮你快速定位谁在“占着茅坑不拉屎”。
3.2 CUDA缓存分配器：PyTorch的“显存管家”
好，现在Storage被标记为可回收了。但显存真的还给系统了吗？
并没有。
PyTorch底层用的是CUDA的缓存分配器（Caching Allocator）。它的工作逻辑很简单：为了省事，把释放的显存先“囤”起来。
你想想看，如果每次分配显存都去调用 cudaMalloc，那性能得多差？系统调用是有开销的。所以PyTorch搞了个“缓存池”：
- 分配时：先从缓存池里找大小合适的块，找不到才去问GPU要
- 释放时：不还给GPU，而是放回缓存池，下次再用
这就是为什么你经常看到 nvidia-smi 显示的显存占用居高不下，但PyTorch内部其实已经“释放”了很多。说白了，nvidia-smi看到的是“已分配+已缓存”，不是“真正在用”。
注意： 缓存分配器有个“坑”——它倾向于分配大块内存。如果你频繁创建和销毁小Tensor，缓存池里会留下很多“碎片”，导致后续大Tensor分配失败。这就是显存碎片化的根源之一。
3.3 显存碎片化：看不见的“内存杀手”
碎片化问题，我估计很多同学都遇到过。场景是这样的：
- 训练刚开始，显存还剩2GB
- 突然要分配一个1.5GB的Tensor
- 结果报OOM了！
为什么会这样？因为显存被切成了很多小块，没有一块连续空间能放下1.5GB。就像你钱包里有一堆硬币，但买瓶水需要一张整钞——你有钱，但花不出去。
碎片化主要来自两个地方：
| 碎片类型 | 产生原因 | 典型场景 |
|---|---|---|
| 内部碎片 | 缓存分配器按2的幂次分配（512B、1KB、2KB...） | 你申请600B，实际拿到1KB，浪费400B |
| 外部碎片 | 频繁分配/释放不同大小的Tensor | 动态网络结构、变长序列处理 |
避坑指南： 我曾经在一个NLP项目中，因为用了动态batch（每个batch的序列长度不一样），导致显存碎片化严重。训练到第100个batch时，明明总显存还剩3GB，但就是分配不出一个200MB的Tensor。后来我改成固定长度padding + 按长度排序，碎片化问题大幅缓解。
3.4 可视化：PyTorch显存分配流程
为了让你更直观地理解整个过程，我画了一张流程图：
3.5 实战建议：如何应对碎片化
知道了原理，咱们得有点“武器”才行。这里分享几个我常用的方法：
方法一：使用 torch.cuda.empty_cache()
这个命令会清空缓存池，把显存还给GPU。但注意，它只释放“缓存”，不释放“正在使用”的Tensor。我一般在每个epoch结束后调用一次，效果不错。
for epoch in range(num_epochs):
train_one_epoch()
torch.cuda.empty_cache() # 清空缓存，减少碎片方法二：设置 PYTORCH_CUDA_ALLOC_CONF
PyTorch 1.9+ 支持通过环境变量控制分配器行为。我最常用的是 max_split_size_mb：
# 限制缓存块最大分割大小，减少碎片
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128这个参数的意思是：缓存池里的块，如果大于128MB就不再分割。能有效减少外部碎片。
方法三：避免动态形状
如果你处理的是变长序列，尽量用padding固定长度。或者像我之前做的，按长度排序后分组。动态形状是碎片化的头号元凶。
嗯，关于PyTorch显存机制，今天就聊到这儿。记住一句话：显存不是不够用，而是不会用。理解了Tensor的Storage机制、缓存分配器的“囤货”逻辑、以及碎片化的成因，你就能从根源上解决OOM问题。
想要更多资料，联系可可