2. YOLO模型内存足迹分析

做NPU推理优化,第一步就是搞清楚内存都花在哪了。我见过不少团队,模型精度调得挺好,一上板子就崩,十有八九是内存没算明白。

YOLO的内存占用,说白了就三大块:权重参数、特征图、中间计算结果。咱们一个一个拆开看。

2.1 权重参数的内存占用

权重参数是模型本身的“骨架”。YOLOv8n大概有300万个参数,YOLOv8x能到6800万。每个参数如果是FP32就是4字节,INT8就是1字节。

模型版本 参数量 FP32 (MB) INT8 (MB)
YOLOv8n 3.0M 11.4 2.9
YOLOv8s 11.1M 42.3 10.6
YOLOv8m 25.9M 98.8 24.7
YOLOv8l 43.7M 166.7 41.7
YOLOv8x 68.2M 260.2 65.1
我的习惯:项目初期直接用INT8估算权重内存,留出余量。如果INT8都超了,赶紧换小模型或者做剪枝,别等到联调再改。

权重参数在推理时是只读的。NPU通常会把它们放在DDR的固定区域,或者直接锁在SRAM里。我个人建议,如果SRAM够用,优先把权重放SRAM——省带宽,省电。

2.2 特征图的内存占用

特征图才是真正的“内存大户”。你想想看,一张640x640的图,经过Backbone、Neck、Head,每一层都会产生特征图。而且这些图在反向传播之前都不能丢——虽然推理不需要反向传播,但NPU的硬件流水线设计决定了,某些中间特征必须保留到当前层计算完。

特征图大小取决于三个因素:

  • 空间尺寸:H x W,从640x640一路下采样到20x20
  • 通道数:从3通道一路升到512甚至1024
  • 数据类型:FP32还是INT8

举个例子,YOLOv8s在Neck阶段的一个特征图,尺寸是80x80x256,INT8下就是80x80x256x1 = 1.6MB。看着不大?但同一时间可能有5-6个这样的特征图活着,加起来就10MB了。

我曾经踩过的坑:有一次在某个NPU芯片上部署YOLOv5s,算力够,带宽够,就是跑不起来。查了两天,发现是特征图总内存超过了芯片的局部SRAM上限,导致频繁DMA搬运,性能直接腰斩。后来做了特征图复用优化才解决。

2.3 中间计算结果的内存占用

中间计算结果,指的是卷积、归一化、激活函数这些算子产生的临时数据。比如:

  • 卷积的im2col缓冲区:某些NPU实现需要把输入展开成矩阵,这个缓冲区可能比特征图本身还大
  • BatchNorm的均值和方差:推理时是常量,但需要临时加载
  • SiLU/ReLU的掩码:某些量化方案需要保留激活掩码
  • 残差连接的加法缓冲区:需要同时保留两条路径的输出

这部分内存比较隐蔽,很多人在估算时容易漏掉。我一般会在模型里挑几个关键层,手动算一下算子需要的临时缓冲区大小,然后加个20%的余量。

2.4 整体内存分布图

下面这张图,是我根据YOLOv8s在INT8下的典型配置画的。你可以直观看到内存都去哪了。

YOLOv8s INT8 推理内存分布 (640x640输入) 权重参数 10.6 MB 只读,建议锁SRAM 特征图 18.5 MB 动态分配,峰值占用 中间计算结果 4.2 MB 算子临时缓冲区 总内存占用 ≈ 33.3 MB 权重 + 特征图 + 中间结果 关键结论: 1. 特征图是最大消耗者,占比超过55% 2. 权重参数固定,优化空间有限 3. 中间结果容易被忽略,但可能成为瓶颈

2.5 不同精度的内存对比

精度选择直接影响内存占用。我整理了一个对比表,方便你快速评估:

精度 权重 (MB) 特征图 (MB) 中间结果 (MB) 总计 (MB) 典型场景
FP32 42.3 74.0 16.8 133.1 服务器/开发板
FP16 21.2 37.0 8.4 66.6 高端NPU
INT8 10.6 18.5 4.2 33.3 边缘NPU首选
INT4 5.3 9.3 2.1 16.7 超低功耗设备

我的建议:边缘NPU部署,直接上INT8。如果芯片SRAM小于32MB,考虑INT4或者模型剪枝。别在FP32上浪费时间——精度提升有限,内存代价太大。

2.6 避坑指南

最后分享几个我实际项目中遇到的坑:

  • 特征图峰值计算:别只看平均,要看同一时间最多有多少个特征图活着。YOLO的Neck部分经常有多个分支并行,峰值可能比平均值高30%
  • DMA对齐开销:NPU的DMA传输通常要求地址对齐(比如32字节对齐),不对齐的话会多分配内存。我见过一个项目,因为这个多占了5%的内存
  • 量化校准表:INT8量化需要校准表,这个表本身也要占内存。虽然不大(几十KB),但别忘了算进去