3、PyTorch内存管理：Tensor存储机制、计算图与中间变量、梯度累积与显存占用
聊到PyTorch的显存管理，我脑子里第一个蹦出来的词就是「玄学」。很多同学跑着跑着模型，突然就OOM了，然后一脸懵地问我：「老师，我明明只用了几个Tensor，怎么就爆显存了？」
嗯，这里面的门道其实不少。PyTorch的显存管理，说白了就是三件事：Tensor怎么存、计算图怎么建、梯度怎么攒。搞懂这三样，你就能从「玄学」变成「科学」。
3.1 Tensor的存储机制：你以为的「删除」可能只是假象
先说说Tensor。很多人以为，我del掉一个Tensor，显存就释放了。其实没那么简单。
PyTorch的Tensor存储分两层：元数据和数据区。元数据就是shape、dtype这些信息，很小。数据区才是真正占显存的大头。你del掉一个Tensor，Python的引用计数减1，但数据区不一定立刻释放——它得等PyTorch的缓存分配器回收。
核心概念：PyTorch有一个自己的内存缓存池。你释放的Tensor数据，先回到这个池子里，下次有同样大小的Tensor申请时，直接复用。这能减少cudaMalloc的调用次数，但也会让你觉得「显存没降下来」。
举个例子：
import torch
# 创建一个1GB的Tensor
x = torch.randn(256, 1024, 1024, device='cuda')
print(f"占用显存: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
# 删除它
del x
print(f"删除后显存: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
# 你会发现，显存没变！为什么会这样？因为PyTorch的缓存分配器还没把这块内存还给驱动。想真正释放？得调用torch.cuda.empty_cache()。但我个人建议，别频繁调用这个函数——它会让缓存池失效，反而拖慢训练速度。
我的经验：在项目中，我一般只在模型评估或切换数据集时调用一次empty_cache。训练过程中，让它自己管理就好。你想想看，频繁清缓存就像每次用完筷子就扔，下次吃饭还得重新买，多浪费。
3.2 计算图与中间变量：显存消耗的隐形杀手
接下来是计算图。这是PyTorch最强大的特性，也是最容易踩坑的地方。
每次你执行一个前向传播，PyTorch都会动态构建一个计算图。这个图里，每个节点是一个操作，每条边是Tensor的依赖关系。为了反向传播，PyTorch必须保留所有中间变量的值——这就是显存消耗的大头。
我画了一张图，帮你理解这个过程：
你看，一个简单的线性层加激活函数，就产生了两个中间变量。如果你的网络有100层，那中间变量就是200个。每个都占显存，加起来就恐怖了。
避坑指南：我曾经在训练一个50层的ResNet时，发现显存占用比理论值大了3倍。排查了半天，才发现是中间变量太多。后来用torch.no_grad()包裹了推理部分，显存立刻降下来了。
那怎么减少中间变量的显存占用？有几个实用技巧：
- 使用in-place操作：比如
relu_()代替relu()，直接在原Tensor上修改，不产生新变量。但要注意，in-place操作会破坏计算图，慎用。 - 梯度检查点（Checkpointing）：用
torch.utils.checkpoint，牺牲一点计算时间，换取显存。原理是反向传播时重新计算中间变量，而不是保存它们。 - 混合精度训练：用float16代替float32，显存直接减半。我现在的项目基本都标配AMP了。
3.3 梯度累积与显存占用：小batch也能跑大模型
最后聊聊梯度累积。这是个很实用的技巧，尤其当你显存不够，但又想用大batch size的时候。
梯度累积的原理很简单：不更新参数，只累积梯度。比如你本来想用batch size=64，但显存只够跑batch size=16。那就跑4次前向+反向，每次梯度累加，第4次再更新参数。
但这里有个坑：梯度本身也占显存！每个参数的梯度都是一个Tensor，大小和参数一样。如果你有1亿参数，梯度就占400MB（float32）。
| 参数规模 | 参数显存（float32） | 梯度显存（float32） | 优化器状态（Adam） | 总计 |
|---|---|---|---|---|
| 1亿 | 400 MB | 400 MB | 800 MB | 1.6 GB |
| 10亿 | 4 GB | 4 GB | 8 GB | 16 GB |
| 100亿 | 40 GB | 40 GB | 80 GB | 160 GB |
你看，梯度占的显存和参数一样多。优化器状态（比如Adam的动量和方差）更是翻倍。所以梯度累积时，你不仅要考虑中间变量，还得算上梯度的开销。
我的习惯：做梯度累积时，我会先用torch.cuda.memory_summary()看一眼当前显存分配情况。这个函数会列出所有Tensor的大小和位置，帮你找到「显存黑洞」。
代码示例：
# 梯度累积的正确姿势
accumulation_steps = 4
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
outputs = model(inputs)
loss = criterion(outputs, labels)
# 除以累积步数，保持loss尺度一致
loss = loss / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad() # 别忘了清零梯度这里有个细节：loss = loss / accumulation_steps。为什么要除？因为每次backward都会累加梯度，如果不除，等效的batch size虽然大了，但学习率也得跟着调，很麻烦。除一下，梯度尺度就一致了。
我曾经踩过的坑：有一次做梯度累积，忘了在optimizer.step()之后调用zero_grad()。结果梯度一直累加，越变越大，模型直接炸了。嗯，从那以后我每次写梯度累积，都会在step后面立刻写上zero_grad，形成肌肉记忆。
3.4 总结：显存管理的三个核心原则
好了，说了这么多，总结一下：
- Tensor存储：别被del骗了，显存释放有延迟。用
torch.cuda.empty_cache()但别滥用。 - 计算图：中间变量是显存消耗的大头。用checkpointing、in-place操作、混合精度来优化。
- 梯度累积：小batch也能跑大模型，但别忘了梯度本身也占显存。记得除累积步数，记得清零梯度。
你想想看，搞懂了这三样，是不是感觉显存管理没那么玄学了？下次再遇到OOM，至少知道该从哪里下手排查了。
想要更多资料，联系可可