2. PyTorch显存机制：Tensor存储原理、CUDA缓存分配器、显存碎片化问题、torch.cuda基础API
好，咱们直接进入正题。显存优化这事儿，说白了就是搞清楚「你的数据在GPU里到底是怎么待着的」。很多同学一上来就调batch size、换模型，结果显存还是炸——为什么？因为你根本没摸透PyTorch的显存脾气。
我个人习惯，先理解底层机制，再动手调优。这一章咱们就掰开揉碎，把PyTorch的显存机制讲明白。
2.1 Tensor存储原理：数据在GPU里长什么样？
先问个问题：你在GPU上创建一个Tensor，它占了多少显存？
嗯，很多人会答「float32就是4字节乘以元素个数」。对，但不全对。Tensor在显存里其实分两块：数据区和元信息区。
- 数据区：存真正的数值，连续内存块。
- 元信息区：存shape、dtype、stride、device等，很小，几十字节。
真正吃显存的是数据区。但有个坑——视图（View）。比如你用tensor[::2]切片，它不复制数据，只是改了stride和offset。元信息变了，数据区还是原来那块。我在项目中遇到过有人疯狂切片，以为没占新显存，结果反向传播时梯度累积，显存直接爆了。为什么？因为autograd会保留中间结果的引用，视图引用的原始数据区一直活着。
核心结论：Tensor占用的显存 = 数据区大小 + 元信息。视图不占新数据区，但可能阻止原始数据释放。
另外，PyTorch的Tensor存储是行优先的。你创建一个(3, 4)的矩阵，它在显存里是连续排成12个数的。但如果你转置一下，stride变了，数据还是那12个数，只是访问顺序变了。这就是为什么转置不占新显存——它只是换了个「看数据的方式」。
2.2 CUDA缓存分配器：PyTorch为什么「赖着显存不还」？
你肯定遇到过这种情况：明明把Tensor删了，torch.cuda.memory_allocated()显示显存没降。你心想「我代码写错了？内存泄漏了？」
别急，这不是bug，是特性。PyTorch底层用了CUDA缓存分配器（Caching Allocator）。它的逻辑很简单：你释放一个Tensor，它不立刻把显存还给操作系统，而是留在自己的缓存池里。下次你再申请显存，直接从池子里拿，省去了和操作系统打交道的开销。
说白了，就是「我帮你管着，别老去烦CUDA驱动」。这能大幅提升训练速度，尤其是频繁创建销毁小Tensor的场景。
但副作用也很明显：你看着nvidia-smi里显存占用居高不下，心里发毛。我曾经调试一个模型，显存占用8GB，我删了所有变量，还是8GB。当时差点重装驱动。后来发现是缓存分配器在作祟。
小技巧：想清空缓存，用 torch.cuda.empty_cache()。它会释放所有未使用的缓存块，还给CUDA驱动。但注意，这不会释放正在被使用的Tensor。
缓存分配器内部维护了一个块列表，每个块有大小、是否空闲等状态。申请显存时，它先找大小匹配的空闲块；找不到就向CUDA申请新块。释放时，块标记为空闲，但物理内存不归还。
这里有个关键参数：PYTORCH_CUDA_ALLOC_CONF。你可以设置max_split_size_mb来控制缓存块的最大分割粒度。我建议，如果你的模型显存波动大，把这个值设小一点，比如max_split_size_mb:32，能减少碎片化。
2.3 显存碎片化问题：为什么显存还有200MB，却申请不了100MB的Tensor？
这个问题很经典。你看着torch.cuda.memory_reserved()显示还剩2GB，但一申请个512MB的Tensor就OOM。为什么？
因为显存碎片化了。缓存分配器手里有很多小块，但连续的大块不够。就像你钱包里有一堆硬币，但买瓶水需要一张整钞——你有钱，但花不出去。
碎片化是怎么产生的？典型场景：你反复创建不同大小的Tensor。比如先创建一个1024MB的，释放；再创建一个512MB的，释放；再创建一个256MB的……每次释放后，缓存分配器把块标记为空闲，但这些块大小不一，位置不连续。下次你要一个768MB的连续块，它找遍所有空闲块，拼不出一个连续的768MB。
避坑指南：我曾经在训练一个多尺度模型时，每步都创建不同分辨率的特征图，结果显存碎片化严重，训练到一半就OOM。后来改成预分配固定大小的缓冲区，问题解决。
怎么缓解碎片化？几个实用方法：
- 预分配大块：训练前先创建几个大Tensor占位，用完再复用。
- 统一Tensor大小：尽量让batch size、特征图尺寸保持一致。
- 用
torch.cuda.memory_summary()查看碎片情况：它会列出所有缓存块的状态。
下面这张图展示了显存碎片化的形成过程：
2.4 torch.cuda基础API：你该记住的几个函数
好了，理论讲完，来点实在的。这些API我几乎每天都会用，你最好也记住。
| API | 作用 | 我的使用场景 |
|---|---|---|
torch.cuda.memory_allocated() |
当前Tensor实际占用的显存（字节） | 监控模型真正用了多少，排除缓存干扰 |
torch.cuda.memory_reserved() |
缓存分配器从CUDA申请的总显存（字节） | 看PyTorch「占着茅坑不拉屎」的量 |
torch.cuda.max_memory_allocated() |
历史峰值显存占用 | 排查OOM时，看峰值出现在哪一步 |
torch.cuda.empty_cache() |
清空缓存分配器中的空闲块 | 切换模型或数据集时，手动释放缓存 |
torch.cuda.memory_summary() |
打印详细的显存分配情况 | 深度调试碎片化问题，看每个块的状态 |
torch.cuda.reset_peak_memory_stats() |
重置峰值统计 | 想重新测量某段代码的峰值显存时 |
举个实际例子。我调试显存泄漏时，经常这么写：
import torch
# 记录初始状态
torch.cuda.reset_peak_memory_stats()
start_allocated = torch.cuda.memory_allocated()
# 你的模型或操作
model = MyModel().cuda()
output = model(input_tensor)
# 查看峰值
peak = torch.cuda.max_memory_allocated()
print(f"峰值显存: {peak / 1024**2:.2f} MB")
# 清理后看是否泄漏
del model, output
torch.cuda.empty_cache()
end_allocated = torch.cuda.memory_allocated()
print(f"泄漏? {(end_allocated - start_allocated) / 1024**2:.2f} MB")
个人习惯：我一般在训练脚本的每个epoch结束时，调用一次torch.cuda.empty_cache()和torch.cuda.reset_peak_memory_stats()。这样每个epoch的显存统计都是干净的，方便对比不同配置的效果。
还有一个冷门但好用的：torch.cuda.set_per_process_memory_fraction(fraction)。它可以限制当前进程能用的显存比例。比如你在一张卡上跑两个实验，各占50%，互不干扰。我在服务器上多人共用GPU时经常用这招。
嗯，这一章的内容差不多就这些。记住：Tensor存储、缓存分配器、碎片化、基础API，这四个点串起来，你就掌握了PyTorch显存机制的骨架。下一章咱们会深入实战，看看怎么用这些知识去优化一个真实模型。