基础概念：张量生命周期、内存分配器、引用计数

各位同学，今天我们来聊聊显存复用最底层的三个概念。说实话，这三个东西搞不清楚，后面讲的各种复用策略你听着都会觉得像空中楼阁。我自己带团队的时候，发现很多新人一上来就研究什么「内存池」「缓存复用」，结果连张量什么时候生什么时候死都说不清楚——那肯定是要出问题的。

张量生命周期：从生到死

每个张量，说白了就是一个数据对象。它在显存里待多久，完全取决于它的生命周期。我个人习惯把生命周期分成四个阶段：

• 创建：张量被实例化，显存被分配
• 使用：参与前向或反向计算
• 引用：被其他操作或变量持有
• 销毁：引用归零，显存被回收

举个例子，你写一行 a = torch.randn(1024, 1024)，这时候张量就诞生了。它在显存里占 4MB（如果是 float32）。然后你拿它做矩阵乘法，算完之后如果没人再引用它，它就死了。嗯，这里要注意——「死」不代表显存立刻还给系统，这涉及到内存分配器的策略，我们后面细说。

我在项目中遇到过一种典型问题：训练一个很大的模型，显存莫名其妙爆了。查了半天发现，有个中间张量在循环里被反复创建，但每次都没及时释放。你想想看，一个循环跑 1000 次，每次多占 100MB，最后就炸了。这就是生命周期管理没做好。

内存分配器：谁在管显存

显存分配器，说白了就是操作系统内存管理在 GPU 上的翻版。但它有个很大的不同——GPU 显存的分配和释放开销非常大。我曾经做过测试，一次 cudaMalloc 调用大概要几十微秒，而一次 cudaFree 甚至可能阻塞整个 CUDA 流。

所以现代深度学习框架（PyTorch、TensorFlow）都用了自己的内存分配器。PyTorch 用的是 CUDA Caching Allocator，它的核心思想就一句话：不急着还。

具体来说：

• 当你释放一个张量，分配器不会立刻调用 cudaFree
• 而是把这个显存块标记为「空闲」，放进缓存池
• 下次有新的张量申请，先看缓存池里有没有大小合适的块
• 有就直接复用，没有才去真的分配新显存

这个机制的好处很明显——减少了 cudaMalloc 的调用次数。我见过一个极端案例，某个模型训练时如果不做缓存，每秒要调用上千次 cudaMalloc，光分配开销就占了 30% 的训练时间。用了缓存分配器之后，这个开销几乎降到了零。

但缓存分配器也有坑。我曾经遇到一个情况：显存明明还有 2GB，但模型就是跑不起来。为什么？因为缓存池里全是碎片——一堆 100MB、200MB 的小块，凑不出一个 1GB 的连续块。这就是显存碎片化问题，后面章节我们会专门讲怎么处理。

引用计数：谁在牵着张量

引用计数，这个概念从 Python 那边继承过来的。每个张量内部维护一个计数器，记录有多少个变量在引用它。当计数器归零，张量就进入销毁流程。

看个简单例子：

import torch

a = torch.randn(100, 100)  # 引用计数 = 1
b = a                      # 引用计数 = 2
c = a                      # 引用计数 = 3

del b                      # 引用计数 = 2
del c                      # 引用计数 = 1
del a                      # 引用计数 = 0 → 销毁

嗯，这里要注意——引用计数归零后，张量占用的显存块会被标记为「可回收」。但正如前面说的，分配器不一定立刻还给系统，而是放进缓存池。

还有一个常见问题：梯度张量的引用。在 PyTorch 中，如果你设置了 requires_grad=True，计算图会持有中间张量的引用，直到反向传播完成。这意味着：

• 前向过程中创建的中间张量，在反向传播之前不会被释放
• 如果你在反向传播之前手动 del 了某个中间张量，计算图会报错
• 反向传播完成后，计算图释放引用，这些张量才真正被回收

所以很多人问：为什么我的显存在前向过程中一直涨，反向之后才降下来？答案就在这里——引用计数被计算图牵着呢。

三者的关系：一张图说清楚

下面我用一张 SVG 图来展示这三个概念是怎么协同工作的：

这张图把整个流程串起来了。你从左上角开始看：张量创建后，引用计数为 1。随着被其他变量引用，计数增加。参与计算时，计数保持不变。当所有引用都释放了，计数归零，张量进入销毁流程。但注意——它没有直接还给系统，而是进了 CUDA Caching Allocator 的缓存池。下次有新张量要分配，优先从池子里拿。只有池子满了，才会真的调用 cudaFree 把显存还给系统。

实际项目中的经验

最后分享一个我自己的教训。有一次我在优化一个推荐系统的训练流程，发现显存占用一直在涨，但明明每个 batch 结束后该释放的都释放了。用 torch.cuda.memory_summary() 一看，好家伙——缓存池里堆了 8GB 的碎片。

原因是什么？每个 batch 生成的张量大小都不一样。第一个 batch 生成 100MB 的，第二个生成 120MB 的，第三个生成 80MB 的……分配器每次都在池子里找，找不到就新分配，但释放时又放回池子。久而久之，池子里全是大小不一的碎片，谁也复用不了谁。

解决办法有两个：

• 尽量让张量大小对齐，比如用 padding 让所有张量都是 128MB 的倍数
• 或者手动清空缓存池：torch.cuda.empty_cache()，但这会触发 cudaFree，有性能开销

我个人更推荐第一种——从源头避免碎片化。后面章节我们会深入讲各种复用策略，但今天这三个基础概念，是你理解所有策略的基石。搞懂了张量怎么生、怎么死、谁在管它的「后事」，后面的内容你就不会觉得乱了。

想要更多资料，联系可可