4、Tensor显存占用分析：从存储格式到实战计算
做深度学习的人，十有八九都遇到过OOM（Out of Memory）。
模型跑着跑着，突然弹出一个显存不足的错误。这时候你怎么办？
我刚开始也只会调小batch size，后来发现——其实很多显存浪费，都是因为我们对Tensor的存储机制不够了解。今天我就带你彻底搞懂，一个Tensor到底是怎么吃掉显存的。
4.1 Tensor的存储格式：dtype、shape、stride
一个Tensor在显存里长什么样？
说白了，就是一块连续的内存区域，加上一些元数据来描述它。
4.1.1 dtype：每个元素占多少字节
这是最直接的显存影响因素。我见过不少新手，上来就用float64，其实完全没必要。
| 数据类型 | 字节数 | 常用场景 |
|---|---|---|
| float32 | 4 | 默认精度，大部分模型用这个 |
| float16 | 2 | 混合精度训练，显存减半 |
| bfloat16 | 2 | Google推荐，动态范围同float32 |
| int8 | 1 | 量化推理，极致省显存 |
| int32 | 4 | 索引、标签等 |
4.1.2 shape：维度决定总元素数
shape就是每个维度的大小。比如一个形状为[32, 3, 224, 224]的Tensor，总共有32×3×224×224 = 4,816,128个元素。
嗯，这里要注意：shape本身不占多少显存，但它决定了后续所有计算的内存需求。
4.1.3 stride：内存布局的隐藏变量
stride是什么？它告诉你在内存中，从一个元素跳到下一个元素需要跨过多少步。
举个例子：
import torch
x = torch.randn(3, 4)
print(x.stride()) # 输出 (4, 1)
# 意思是：行之间跨4个元素，列之间跨1个元素为什么stride重要？
因为有些操作（比如转置、切片）不会复制数据，只是修改stride。这在节省显存的同时，也可能带来一些坑。
4.2 显存计算公式：一个Tensor到底占多少
公式其实很简单：
用代码表示就是：
def tensor_memory(tensor):
return tensor.numel() * tensor.element_size()
# 例子
x = torch.randn(1024, 1024, dtype=torch.float32)
print(f"占用显存: {tensor_memory(x) / 1024 / 1024:.2f} MB")
# 输出: 占用显存: 4.00 MB你想想看，一个1024×1024的矩阵，float32，正好4MB。如果换成float16，就只有2MB。
但实际训练中，显存占用远不止这些。为什么？因为还有梯度、优化器状态、中间激活值。
4.3 梯度与优化器状态的显存开销
这部分是很多人容易忽略的。我刚开始做分布式训练时，就因为这个吃了大亏。
4.3.1 梯度：每个参数都有一个对应的梯度
梯度Tensor和参数Tensor形状完全一样，dtype也相同。所以：
- 参数占多少显存，梯度就占多少
- 也就是说，光梯度和参数，就已经翻倍了
4.3.2 优化器状态：Adam最吃显存
不同的优化器，保存的状态不一样：
| 优化器 | 额外状态 | 显存开销 |
|---|---|---|
| SGD | 无（或momentum buffer） | 0 ~ 1倍参数大小 |
| Adam | 一阶动量、二阶动量 | 2倍参数大小 |
| AdamW | 同Adam | 2倍参数大小 |
- 参数本身：400MB
- 梯度：400MB
- Adam状态：800MB
- 合计：1.6GB
所以，我建议你：如果显存紧张，可以考虑用SGD或者LAMB优化器。Adam虽然好用，但代价确实不小。
4.4 中间激活值的显存占用：真正的显存大户
这才是训练时显存爆炸的元凶。
什么是中间激活值？就是前向传播过程中，每一层的输出。反向传播计算梯度时，需要用到这些值。
4.4.1 激活值怎么算
对于每一层：
- 输入激活值：batch_size × 输入通道 × 高 × 宽
- 输出激活值：batch_size × 输出通道 × 高 × 宽
- 如果是Transformer，还有attention矩阵、key/value缓存等
举个例子，一个简单的CNN：
# 假设输入是 [32, 3, 224, 224]
# 经过一个卷积层，输出 [32, 64, 112, 112]
# 中间激活值大小：32 * 64 * 112 * 112 * 4 = 约102MB
# 如果网络有50层，光激活值就可能到5GB以上4.4.2 如何减少激活值显存
有几个常用方法：
- 梯度检查点（Gradient Checkpointing）：不保存所有激活值，反向传播时重新计算。时间换空间，显存可以降低到原来的1/3甚至更低。
- 使用更小的batch size：最直接，但会影响训练效率。
- 混合精度训练：激活值用float16存储，显存直接减半。
4.5 知识体系总览
下面这张图，帮你把整个Tensor显存占用的逻辑串起来：
这张图把整个脉络理清楚了。你从存储格式出发，算清楚每个Tensor占多少，再把梯度、优化器、激活值都加上去，就能准确预估训练时的显存需求。
我个人习惯是：在写训练脚本之前，先手动算一遍显存。如果发现超了，就提前做优化，而不是等到跑起来才报错。
好了，这一章的内容就到这里。记住：显存不是无限的，但理解它的分配规律，能让你在有限的资源里跑更大的模型。