4、内存优化基础：显存与内存模型、内存碎片化问题、池化与复用策略
说到YOLO部署，很多人第一反应是「模型精度掉没掉」「推理速度快不快」。
但我得说，内存这块才是真正容易翻车的地方。你想想看，模型在GPU上跑着跑着突然OOM，或者显存占用忽高忽低，这种问题我遇到过太多次了。今天咱们就把内存优化这块地基打牢。
4.1 显存与内存模型：你得知道数据住哪儿
先搞清楚一个基本概念。YOLO部署时，数据主要在三个地方流动：CPU内存、GPU显存、还有CPU与GPU之间的传输通道。
我个人习惯把显存想象成一个「高速工作台」。模型推理时，权重、输入图像、中间特征图、输出结果，全得摆在这个台子上。而CPU内存呢，更像是「仓库」，负责存原始数据、预处理结果、后处理结果。
这里有个常见的误区——很多人以为显存只存模型权重。其实不是的。我举个例子，YOLOv8输入一张1080p的图像，光输入张量就要占掉大约6MB（640x640x3x4字节）。如果batch size是32，那就是192MB。再加上每一层的中间特征图，你算算看，一个中等规模的模型，显存占用轻松上GB。
显存占用核心公式：
总显存 ≈ 模型权重 + 输入张量 + 所有中间特征图 + 输出张量 + 框架开销
我在项目中遇到过最离谱的一次，是有人把batch size设成64，结果显存直接爆了。后来一查，中间特征图占了总显存的60%以上。嗯，这就是典型的「只算权重没算特征图」的坑。
4.2 内存碎片化问题：看不见的性能杀手
内存碎片化，说白了就是显存被切得七零八落，像一块被狗啃过的蛋糕。
为什么会这样？
YOLO推理过程中，会频繁地分配和释放各种大小的张量。比如前向传播时，每一层都要分配新的特征图，用完又释放。这种「频繁申请-释放」的模式，时间一长，显存里就会出现很多小空隙。
这些小空隙单独看都不大，但加起来可能占掉几百MB。更气人的是，当你需要一块连续的大显存时（比如分配一个大的输出张量），明明总剩余显存够用，但因为不连续，就是分配不出来——然后程序就OOM了。
避坑指南：
我曾经在一个边缘设备上部署YOLOv5s，跑了大概2小时后突然报OOM。排查了半天，不是内存泄漏，就是碎片化问题。重启进程后一切正常，但跑几小时又复现。最后用显存分析工具一看，碎片率高达35%。
怎么检测碎片化？我一般用两个方法：
- nvidia-smi 看显存占用：如果显存占用持续增长但推理速度反而下降，大概率是碎片化
- 框架自带工具：PyTorch可以用
torch.cuda.memory_summary()查看详细分配情况
碎片化的直接后果就是推理延迟抖动。你想想看，本来一次推理只要10ms，结果因为碎片化，显存分配花了50ms，这谁受得了？
4.3 池化与复用策略：把显存当共享单车骑
解决碎片化问题，最有效的思路就是「池化与复用」。说白了，就是别用完就扔，而是回收再利用。
我常用的策略有三种：
4.3.1 显存池（Memory Pool）
显存池的原理很简单：预先申请一大块显存，然后自己管理内部的分配和回收。不再依赖框架的默认分配器。
举个例子，TensorRT就内置了显存池机制。你可以通过设置 IGpuAllocator 来接管显存管理。我自己在项目中写过一套简单的显存池，核心逻辑就这几行：
class MemoryPool:
def __init__(self, total_size):
self.pool = cuda.mem_alloc(total_size)
self.free_blocks = [(0, total_size)]
self.used_blocks = {}
def alloc(self, size):
# 从free_blocks中找到合适大小的空闲块
for i, (offset, block_size) in enumerate(self.free_blocks):
if block_size >= size:
# 分配并切分
self.free_blocks.pop(i)
if block_size > size:
self.free_blocks.append((offset + size, block_size - size))
self.used_blocks[offset] = size
return self.pool + offset
raise MemoryError("池子不够用了")
def free(self, offset):
# 回收并合并相邻空闲块
size = self.used_blocks.pop(offset)
self.free_blocks.append((offset, size))
self._merge_blocks() # 合并相邻块，减少碎片个人经验：显存池的大小一般设为模型峰值显存占用的1.2倍。太小了不够用，太大了浪费。我习惯先跑一次profiling，看看峰值是多少，再定池大小。
4.3.2 张量复用（Tensor Reuse）
YOLO推理时，很多中间张量的生命周期是重叠的。比如卷积层的输入和输出，不会同时存在。这时候就可以复用显存。
具体做法是：分析计算图，找出那些「生命周期不重叠」的张量，让它们共享同一块显存。
我记得在优化YOLOv8部署时，通过张量复用，把中间特征图的显存占用从1.2GB降到了600MB。怎么做的？就是发现backbone的某些层和neck的某些层，它们的特征图不会同时被用到，那就共用一个buffer呗。
4.3.3 对象池（Object Pool）
这个主要针对后处理阶段。YOLO的后处理要生成很多检测框、置信度、类别等对象。如果每个对象都new/delete，碎片化会非常严重。
我的做法是：预先分配一个对象池，后处理时从池里取，用完放回去。
class DetectionPool:
def __init__(self, pool_size=1000):
self.pool = [Detection() for _ in range(pool_size)]
self.available = list(range(pool_size))
def get(self):
if not self.available:
# 池子不够了，动态扩容
self._expand()
idx = self.available.pop()
return self.pool[idx]
def release(self, det):
# 重置对象状态
det.reset()
self.available.append(id(det))三种策略对比：
| 策略 | 适用场景 | 效果 | 实现难度 |
|---|---|---|---|
| 显存池 | 整体显存管理 | 减少碎片化，降低分配延迟 | 中等 |
| 张量复用 | 中间特征图 | 大幅降低峰值显存 | 较高 |
| 对象池 | 后处理阶段 | 减少小对象碎片 | 低 |
4.4 实战建议：从哪开始优化？
如果你现在就要动手优化YOLO部署的内存，我建议按这个顺序来：
- 先做profiling：搞清楚显存到底花在哪了。用
torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()看峰值 - 检查碎片化程度：连续跑100次推理，看显存占用是否持续增长
- 从最简单的开始：先实现对象池，这个改动最小，效果立竿见影
- 再考虑张量复用：如果峰值显存还是高，分析计算图，找复用机会
- 最后上显存池：如果碎片化严重，或者框架默认分配器表现不佳，再自己写池
注意：不要一上来就搞显存池。我见过有人花了两周写了个完美的显存池，结果发现框架自带的分配器已经够用了。先测量，再优化，这是铁律。
好了，这一章的内容就到这。内存优化是个细活，但只要你把显存模型、碎片化问题、池化复用这三个点吃透了，后面再讲具体优化技巧时，你就能很快上手。下一章咱们聊聊「推理引擎的内存管理机制」，到时候会结合TensorRT和ONNX Runtime的实际案例来讲。