2. PyTorch显存管理机制：Tensor存储与视图、计算图与中间变量、梯度与优化器状态、显存分配与释放策略
显存管理，说白了就是跟GPU内存打交道。很多同学跑模型时突然OOM（显存溢出），然后一脸懵——明明模型不大啊？其实问题往往出在你对PyTorch的显存管理机制不够了解。我刚开始做训练优化时，也踩过不少坑，今天就把这些经验掰开揉碎了讲给你听。
2.1 Tensor存储与视图：数据到底存在哪？
Tensor是PyTorch的基本单位。但你知道吗？一个Tensor的显存占用，其实分两部分：数据区和元信息区。
- 数据区：存的是真正的数值，比如float32的权重、int64的索引。这部分占大头。
- 元信息区：存的是shape、dtype、stride、device等。很小，但很重要。
我见过有人写代码时，每层都创建一个新的Tensor，结果显存直接爆了。其实很多操作可以共享底层数据——这就是视图（View）的威力。
核心概念：存储（Storage） vs 视图（View）
一个Tensor背后有一个Storage对象，它才是真正占显存的地方。多个Tensor可以共享同一个Storage，只是通过不同的shape、stride来“看”数据。比如：
import torch
x = torch.randn(4, 4) # 分配了16个float32的存储空间
y = x[0:2, 0:2] # 视图，没有新分配显存
z = x.contiguous() # 如果x不是连续的，这里会重新分配
print(x.storage().data_ptr() == y.storage().data_ptr()) # True，共享存储
print(x.storage().data_ptr() == z.storage().data_ptr()) # 可能False嗯，这里要注意：view()、reshape()、切片操作，大部分时候返回的是视图。但如果你对视图做了.contiguous()或者.clone()，那就真的重新分配显存了。我在项目中遇到过有人用.view()之后又调了.contiguous()，白白多占了一倍显存。
我的习惯：能用切片就用切片，能用view就用view。只有在需要修改数据且不影响原Tensor时，才用clone。
2.2 计算图与中间变量：前向传播的“隐形杀手”
PyTorch是动态图框架。每次前向传播，它都会构建一个计算图。这个图里记录了所有操作和中间结果——目的是为了反向传播时算梯度。
你想想看，一个ResNet-50前向传播，中间会产生多少激活值？每一层的输出、每个激活函数的输入输出……这些全部存在显存里。这就是为什么训练时显存占用远大于推理时。
关键点：只有requires_grad=True的Tensor才会被计算图追踪。但中间变量（比如ReLU的输出）也会被保留，直到反向传播结束。
# 举个例子
x = torch.randn(1, 3, 224, 224, requires_grad=True)
conv = torch.nn.Conv2d(3, 64, 3)
y = conv(x) # 中间变量y被保留
z = y.relu() # 中间变量z被保留
loss = z.sum() # loss也是中间变量
# 此时显存里存了：x, conv.weight, conv.bias, y, z, loss
# 反向传播后，y和z会被释放
loss.backward()我曾经调试过一个模型，显存占用异常高。后来发现是某个自定义层里，把中间结果存到了self.xxx里，导致计算图无法释放。嗯，这是个经典坑——不要用成员变量保存中间结果，除非你明确知道自己在做什么。
避坑指南：我曾经在写GAN时，把生成器的中间特征图存到了列表里用于可视化，结果显存直接翻了三倍。后来改用torch.no_grad()或者.detach()才解决。
2.3 梯度与优化器状态：反向传播后的“隐形负担”
反向传播完成后，每个可训练参数都会有一个.grad属性。这个梯度Tensor和参数本身一样大。所以如果你有100M参数，梯度就额外占100M * 4字节 = 400MB（float32）。
但这还没完。优化器还会维护自己的状态。比如：
| 优化器 | 额外状态 | 显存开销（相对参数） |
|---|---|---|
| SGD | 无（momentum=0时） | 0 |
| SGD + momentum | 动量缓冲区 | 1倍 |
| Adam | 一阶矩 + 二阶矩 | 2倍 |
| AdamW | 一阶矩 + 二阶矩 | 2倍 |
说白了，你用Adam训练一个1GB的模型，光优化器状态就要吃掉2GB显存。加上参数本身1GB、梯度1GB、中间变量若干……显存就是这么没的。
我的建议：如果显存紧张，可以试试SGD + momentum，或者用torch.optim.AdamW配合foreach=False减少一些临时缓冲区。另外，gradient checkpointing也能省中间变量，但会慢一些。
2.4 显存分配与释放策略：PyTorch的“内存池”机制
PyTorch不会每次分配Tensor都去问GPU要显存。它内部维护了一个缓存分配器（Caching Allocator）。当你释放一个Tensor时，显存不会立刻还给系统，而是留在缓存池里，供后续分配使用。
这样做的好处是快——避免频繁的cudaMalloc/cudaFree开销。但坏处是：你看到的nvidia-smi显存占用，可能比实际需要的多。因为缓存池里还躺着很多“已释放但未归还”的块。
核心函数：
torch.cuda.empty_cache()：清空缓存池，归还显存给系统。但注意，这不会释放正在使用的Tensor。torch.cuda.memory_summary()：查看详细的显存分配情况，调试时很有用。torch.cuda.reset_peak_memory_stats()：重置峰值显存统计。
# 调试显存的好方法
import torch
model = torch.nn.Linear(1000, 1000).cuda()
optimizer = torch.optim.Adam(model.parameters())
# 记录初始显存
torch.cuda.reset_peak_memory_stats()
torch.cuda.empty_cache()
x = torch.randn(1024, 1000).cuda()
y = model(x)
loss = y.sum()
loss.backward()
optimizer.step()
# 查看峰值显存
print(f"峰值显存: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
print(f"当前显存: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")我习惯在训练循环里定期调用torch.cuda.empty_cache()，尤其是在验证集评估之后。不过别太频繁，否则反而影响性能。
注意：如果你用torch.cuda.empty_cache()后显存还没降下来，说明有Tensor的引用没释放。检查一下是不是有变量不小心被保存到了列表或字典里。
知识体系总览
下面这张图总结了本章的核心逻辑，我建议你保存下来，以后遇到显存问题可以对照着排查。
好了，这一章的内容就到这里。显存管理是训练优化的基本功，理解了这些机制，你就能更精准地定位OOM问题，也能更高效地利用每一MB显存。