内存优化：内存复用、内存池、数据布局优化

各位同学，今天我们来聊聊算子优化里最实在的一个话题——内存优化。说实话，我见过太多团队把精力全放在计算加速上，结果发现瓶颈根本不在算力，而在内存。你想想看，GPU算得再快，数据搬来搬去的时间都够泡杯咖啡了。

我个人习惯把内存优化分成三个层次：内存复用、内存池、数据布局优化。这三层是递进关系，从粗粒度到细粒度，从系统层面到算子层面。咱们一个一个来拆解。

1. 内存复用：别让显存白白浪费

什么叫内存复用？说白了就是一块显存，多个用途。我在项目中遇到过这样一个场景：一个Transformer模型，中间有十几个张量，每个都申请独立显存，结果显存直接爆了。后来一分析，很多张量的生命周期根本不重叠。

举个例子，假设你有两个算子：

• 算子A输出张量X，生命周期是第1~3步
• 算子B输出张量Y，生命周期是第4~6步

那X和Y完全可以共用同一块显存。这就是内存复用的核心思想——通过分析张量的生命周期，让不冲突的张量共享内存。

实现上，我们通常用引用计数或区间分析来做。引用计数简单粗暴：每个张量记录被引用的次数，归零就回收。区间分析更精细：把整个计算图的时间轴画出来，每个张量占一个时间段，然后做区间调度。

2. 内存池：减少分配释放的开销

内存复用解决的是「怎么省空间」，内存池解决的是「怎么省时间」。你想想看，每次调用cudaMalloc分配显存，那都是毫秒级的开销。如果模型里有几百个张量，光分配释放就能吃掉几十毫秒。

内存池的思路很简单：提前申请一大块显存，然后自己管理分配和回收。就像你开公司，与其每次用打印机都去现买，不如一次性囤一批。

常见的实现方式有两种：

• 伙伴系统（Buddy System）：把内存按2的幂次分块，分配时找最合适的块，释放时合并相邻空闲块。优点是碎片少，缺点是内部碎片（比如你要3MB，它给你4MB）。
• Slab分配器：针对固定大小的张量，预分配相同大小的内存块。比如你的模型里全是[1, 3, 224, 224]的输入，那就专门准备一批这个大小的块。

我个人建议：先做内存复用，再做内存池。复用把空间省下来，池化把时间省下来。顺序别搞反了。

3. 数据布局优化：NHWC vs NCHW

这个知识点，说白了就是数据在内存里怎么排。你想想看，同样的数据，排法不同，访存效率天差地别。

先解释一下这两个格式：

• NCHW：Batch、Channel、Height、Width。通道维在最内层，每个通道的数据连续存储。
• NHWC：Batch、Height、Width、Channel。像素维在最内层，每个像素的所有通道连续存储。

举个例子，一张[1, 3, 4, 4]的图片：

• NCHW：先存所有R通道的16个像素，再存G通道，再存B通道。
• NHWC：先存第一个像素的RGB，再存第二个像素的RGB，以此类推。

那到底用哪个？我直接说结论：

场景	推荐格式	原因
卷积操作（GPU）	NCHW	cuDNN对NCHW做了深度优化，访存连续
卷积操作（CPU）	NHWC	CPU缓存友好，SIMD向量化更容易
逐元素操作（ReLU、BN）	NHWC	每个像素的所有通道一起处理，访存局部性好
全连接层	NCHW	矩阵乘法需要连续的内存布局

还有一个细节：格式转换本身也有开销。从NCHW转NHWC需要重排内存，这个操作在GPU上叫permute或transpose。如果频繁转换，得不偿失。我的建议是：在计算图的入口和出口各做一次转换，中间保持统一格式。

4. 知识体系总览

下面这张图，是我自己总结的内存优化知识体系。你可以把它当作一个检查清单：

5. 实战建议

最后，我给大家几个实操建议：

1. 先做profiling：别凭感觉优化。用nvidia-smi或Nsight看看显存占用和分配次数，找到真正的瓶颈。
2. 内存复用是低垂的果实：大多数模型都有30%~50%的复用空间。我见过一个ResNet-50，简单复用后显存从2.3GB降到1.4GB。
3. 内存池别贪大：预分配太多显存，其他进程就没法用了。我一般预留模型峰值需求的1.2倍。
4. 数据布局要统一：整个计算链用同一种格式，避免频繁转换。如果实在要转，用torch.channels_last这类框架内置支持。

好了，内存优化这块就讲到这里。记住一句话：计算是显性的，内存是隐性的。很多时候，把内存管好了，性能自然就上去了。

想要更多资料，联系可可