RK3588 上 YOLOv5s 部署:NMS 后处理性能优化与 C++ 加速实践

📅 发布时间:2026/10/1 15:01:54
RK3588 上 YOLOv5s 部署:NMS 后处理性能优化与 C++ 加速实践
RK3588 上部署 YOLOv5s 这个系列前四篇我分别讲了环境准备、模型导出、RKNN 推理和后处理解码。走到这一步按理说该收工了但真正跑起整条 pipeline 的时候才发现NPU 推理虽然很快整个检测流程却卡在了一个毫不起眼的环节上NMSNon-Maximum Suppression非极大值抑制。YOLOv5s 在 640×640 输入下会产出 25200 个候选框这些框的 NMS 计算完全在 CPU 上进行NPU 一点忙都帮不上。我用 Python 写的第一版后处理单帧要 150ms 左右推理部分再快也白搭。后来我把 NMS 整个换成 C做了一系列裁剪候选量和并行化的事最终把耗时压到了 0.42ms 上下换算下来大概加速了 359 倍。这篇就把我的完整思路和实测数据都摊开讲讲。1. 为什么 YOLOv5s 部署到最后卡在了 CPU 的 NMS 上1.1 RK3588 的算力分布与后处理现实RK3588 是一颗典型的 ARM 平台 SoC四颗 Cortex-A76 大核加四颗 Cortex-A55 小核NPU 算力放在边缘设备里也算能打。但这里有个容易被忽略的问题NPU 擅长的是卷积、激活、池化这类规整算子NMS 这种“排序 循环 条件删除”的非规整逻辑很难塞进 NPU 的算子栈里。官方 RKNN 工具链也不支持把 NMS 直接放进模型图里跑。所以现实中绝大多数 YOLO 部署方案最终都必须把自己的 NMS 放到 CPU 上执行。在 RK3588 上做 CPU 后处理环境比 x86 苛刻不少。A76 大核虽然能到 2.4GHz 左右但跟桌面 CPU 比缓存和分支预测能力都有限更别提当你用 Python 写循环密集逻辑时解释器本身的额外开销会直接把 CPU 性能打骨折。我最初就是在 Python 里写了后处理解码推理端到端大概 40ms 左右听起来还行但把 NMS 加进去后单帧直接突破 190ms这里面的坑不是模型而是后处理。1.2 25200 个候选框是怎么来的YOLOv5s 有三个检测头下采样倍数分别是 8、16、32对应到 640×640 输入就是 80×80、40×40、20×20 三张特征图。每个特征图上的每个位置有 3 个 anchor所以总候选框数量是(80×80 40×40 20×20) × 3 25200每个候选框对应 85 个 float分别是 4 个坐标、1 个 objectness 置信度、80 个类别概率。模型推理输出这一堆东西可能只要十几毫秒可后处理要对这 25200 个框逐框解码、过滤、排序然后进入 NMS 的双层循环。NMS 本质上是计算两两框之间的 IoU最坏情况复杂度是 O(N²)。当 N 是 25200 时这个双重循环的规模是 6 亿多次 IoU任何语言都扛不住。这也是我第一篇博文里就强调过的YOLO 部署的帧率瓶颈往往不在模型推理而在后处理。你 NPU 再快一旦 NMS 拖后腿整体帧率照样拉胯。1.3 Python 后处理慢不是因为你不会写可能有朋友会说用 numpy 向量化 IoU 是不是就快了我最初也这么干结果被现实教育了。NMS 的串行选择过程决定了你没法像卷积那样一次性对所有框做矩阵运算选出一个高置信度框后要基于它去抑制其他框然后继续选下一个这个依赖链是跑不掉的前后串行逻辑。即使用 numpy 算 IoU 矩阵过滤、索引、排序还是得在 Python 层重复操作每一次数组操作都要产生临时对象和类型检查。IoU 本身是个非常简单的算式小学除法水平但被 Python 解释器的对象模型放大之后25200 个框的 NMS 就变成了 150ms 的怪兽。另外如果不按类别区分直接对所有候选框做 NMS会把不同类别的重叠框也互相误删。所以正确做法是按类别分组每个类别内部独立做抑制。这个分组逻辑在 Python 里又得三层循环嵌套。慢是必然的。1.4 现成的 NMS 库为什么也不能直接用部署到 RK3588 后PyTorch 环境一般不会带到板子上torchvision.ops.nms基本不用考虑。OpenCV 里有cv::dnn::NMSBoxes这个函数在简单场景能用但它有几个问题一是它对多类别支持不够灵活你得按类别反复调用二是它没有“每个类别只取 TopK 候选框”的能力三是它的内部实现对多核利用率一般。实际测下来用 OpenCV 的 NMSBoxes 处理 25200 个框也有十几毫秒跟手写 C 优化后的零点几毫秒差距非常大。自己写 NMS 并不是重复造轮子而是为了把后处理的每一毫秒都抠出来在 RK3588 上尤其值得。2. 从 Python 到 C第一版朴素实现先跑通再谈优化2.1 先用最简单的 BBox 结构把流程打通我不喜欢一上来就把代码写得很复杂。第一版 C NMS 就做了最直白的事定义一个结构体存坐标、得分、类别然后排序循环删框。struct BBox { float x1, y1, x2, y2; // 目标框左上/右下坐标 float score; // 最终置信度由 objectness × 类别概率得到 int class_id; // 类别编号NMS 只在同类别内进行 };这里我特意在解码阶段把 YOLOv5s 输出的cx, cy, w, h转成了x1, y1, x2, y2的绝对坐标格式。这样做的好处是 NMS 阶段不需要再引入任何坐标转换IoU 计算只管边界框交集就行。类别的存在则是为了后续按类别分组避免不同类别的框互相误删。2.2 朴素 NMS 的代码实现与隐藏代价第一版朴素 NMS 长下面这样std::vectorBBox nms_naive(std::vectorBBox boxes, float iou_threshold) { std::sort(boxes.begin(), boxes.end(), [](const BBox a, const BBox b) { return a.score b.score; }); std::vectorBBox result; while (!boxes.empty()) { BBox best boxes.front(); result.push_back(best); boxes.erase(boxes.begin()); for (auto it boxes.begin(); it ! boxes.end();) { if (IoU(best, *it) iou_threshold) { it boxes.erase(it); } else { it; } } } return result; }这段代码逻辑清楚但性能隐患极其明显std::vector::erase每次删除元素都会触发后续所有元素前移复杂度 O(N)。在 NMS 循环里反复 erase整体可能退化到 O(N³)。我在 25200 个候选框上测试时这个函数虽然能跑完但哪怕开了-O3单帧也到了 18.6ms。当时内心其实是崩溃的C 比 Python 快了约 8 倍但这个数字根本没法让整个检测流程达到实时。这也印证了那句话流程先跑通只是第一步真正能上生产的优化还没开始。2.3 第一版实测数据为什么只有 8.2 倍加速版本说明单帧耗时Python 循环版双层 for numpy 计算 IoU152.4 msC 朴素版直接 sort erase IoU 循环18.6 ms为什么 C 只比 Python 快了 8 倍而不是几十倍原因是 Python 版没有完全傻循环IoU 计算用了 numpy 向量化省掉了一部分解释器开销C 版本则被erase的元素搬迁和大规模 O(N²) 计算拖住。两者各自都有浪费速度差距没有拉开。另外第一版还有一个正确性问题它没有按类别做隔离。不同类别的两个高度重叠框也会被抑制掉这在某些场景下是严重的逻辑错误。所以第一版存在的意义就是验证接口和流程真正做优化还得从头梳理算法。3. 加速 359 倍的关键不是堆代码是削候选框3.1 先削数据规模而不是先上 SIMDNMS 的时间复杂度里真正决定量级的是输入框数量 N。如果能从 25200 切成几百个候选框计算量按 O(N²) 直接下降 7000 倍左右这是任何语言级优化都做不到的数字。我当时把优化重心放在三条路置信度过滤、类别内 TopK、排序后只保留前 K。YOLOv5s 解码完的候选框里绝大多数框的置信度都低到可以忽略真正有意义的可能就几百个。我们用置信度阈值 0.25 做一次粗筛把明显是背景的框扔掉再在每类内部取分数最高的 300 个框进入 NMS准确率和全量 NMS 几乎一致。别小看这一步。我在 100 张图上做过统计TopK300 与 TopK2000 在标准验证集上的 mAP 差异不到 0.1%漏检基本可以忽略。当你实时性吃紧时这类“算法级降复杂”的操作远比堆 SIMD 划算。3.2 用 partial_sort 而不是全局排序第一版用了std::sort把所有候选框都排好。但 NMS 只需要按分数从高到低一批批取框后面几万个低分框的顺序根本无所谓。C 标准库里std::partial_sort正好干这个它只保证前 K 个元素有序剩余部分乱序时间复杂度是 O(N log K)K 远小于 N 时比全排序快很多。组合置信度过滤和 TopK 选框代码大概长这样std::vectorBBox topk_by_class(const std::vectorBBox boxes, int topk, float conf_threshold) { const int num_classes 80; std::vectorstd::vectorBBox grouped(num_classes); for (const auto b : boxes) { if (b.score conf_threshold) { grouped[b.class_id].push_back(b); } } std::vectorBBox selected; for (auto vec : grouped) { if (vec.empty()) continue; int k std::min(topk, static_castint(vec.size())); std::partial_sort(vec.begin(), vec.begin() k, vec.end(), [](const BBox a, const BBox b) { return a.score b.score; }); selected.insert(selected.end(), vec.begin(), vec.begin() k); } return selected; }这样处理之后进入 NMS 的输入规模从 25200 变成了平均 350 个左右。为什么不是固定 300因为很多类别根本没满 300 个候选只有人和车这类主类别会接近上限。3.3 用标记数组代替 erase并预计算面积削完候选框剩下的就是让 NMS 循环本身不再做无意义的拷贝。核心改动有两个用suppressed布尔数组标记被抑制的框而不是从 vector 中真正删除。每个框的面积只在开头算一次存进areas数组之后 IoU 直接复用避免重复计算宽高乘积。优化后的核心循环大致如下std::vectorBBox nms_fast(const std::vectorBBox boxes, float iou_threshold) { const int n static_castint(boxes.size()); std::vectorfloat areas(n); std::vectoruint8_t suppressed(n, 0); std::vectorBBox result; result.reserve(64); for (int i 0; i n; i) { areas[i] (boxes[i].x2 - boxes[i].x1) * (boxes[i].y2 - boxes[i].y1); } for (int i 0; i n; i) { if (suppressed[i]) continue; const BBox a boxes[i]; result.push_back(a); for (int j i 1; j n; j) { if (suppressed[j]) continue; if (a.class_id ! boxes[j].class_id) continue; const BBox b boxes[j]; float ix1 std::max(a.x1, b.x1); float iy1 std::max(a.y1, b.y1); float ix2 std::min(a.x2, b.x2); float iy2 std::min(a.y2, b.y2); float inter_w std::max(0.0f, ix2 - ix1); float inter_h std::max(0.0f, iy2 - iy1); float inter inter_w * inter_h; float iou inter / (areas[i] areas[j] - inter); if (iou iou_threshold) { suppressed[j] 1; } } } return result; }这里假设boxes已经按类别分组并且在每组内按分数降序排好。实际操作中我会对每个类别分组后的子向量分别调用nms_fast再把结果合并。由于不同类别之间直接通过class_id判断跳过就算类别混在同一个数组里也不会互相误删。3.4 多类别并行4 个 A76 大核这样用YOLOv5 有 80 类NMS 天然按类别独立。框架不会要求跨类别的框互相抑制所以类别之间完全可以并行跑。RK3588 有 4 个 A76 大核和 4 个 A55 小核第一反应可能是开 8 线程跑满但实测下来让 A55 参与 NMS 反而会让耗时抖动更严重因为小核性能太弱而且线程调度本身有开销。我最终固定只开 4 个线程并绑定到 4 个 A76 大核上用 OpenMP 做类别维度的并行std::vectorstd::vectorBBox per_class_results(num_classes); #pragma omp parallel for num_threads(4) schedule(dynamic) for (int c 0; c num_classes; c) { if (!grouped[c].empty()) { per_class_results[c] nms_fast(grouped[c], iou_threshold); } }schedule(dynamic)是为了避免某些类别候选框特别多其它线程空转。这一版跑下来NMS 耗时从单线程的 2.85ms 降到了 0.42ms 左右整个后处理流程的延迟一下子变得可以忽略不计。4. 编译器、NEON 与浮点陷阱RK3588 上 C 加速的最后一公里4.1 交叉编译与编译选项哪些该开哪些值得犹豫RK3588 是 ARM 平台开发机上可以用交叉编译工具链直接编。我用的编译命令大概是aarch64-linux-gnu-g -stdc17 -O3 -marcharmv8.2-afp16 -ffast-math -funroll-loops \ nms.cpp -o nms_test -fopenmp逐项解释一下-O3是基础优化肯定开。-marcharmv8.2-afp16是针对 A76 核心微架构的优化选项。RK3588 的 A76 支持 ARMv8.2 扩展fp16 只是顺手带上如果你代码里没用半精度浮点也可以去掉。-ffast-math允许编译器对浮点运算做更激进的优化比如重排运算顺序。NMS 场景下输入的坐标和置信度都是有限浮点数不是极端数据我开着这个选项做了长时间运行没有碰到 NaN 或精度异常。但如果你的工程对浮点结果要求很严格建议先对比开关该选项后的差异。-funroll-loops对循环密集的 IoU 计算有一点点收益不强求。每次编完记得用ldd检查动态库依赖别让板子上找不到 libgomp 之类的运行时库。也可以用-static-libstdc -static-libgcc减少运行时意外。4.2 NEON 加速 IoU 的尝试我主动放弃了ARM 平台的 NEON 指令一次能算 4 个 float理论上很适合 IoU 这类乘加运算。我试过一个 NEON 版本用vmaxq_f32、vminq_f32、vmulq_f32批量处理坐标思路大概是这样float32x4_t ix1 vmaxq_f32(a_x1, b_x1); float32x4_t iy1 vmaxq_f32(a_y1, b_y1); float32x4_t ix2 vminq_f32(a_x2, b_x2); float32x4_t iy2 vminq_f32(a_y2, b_y2); // 然后用 vmaxq_f32(vsubq_f32(ix2, ix1), vdupq_n_f32(0)) 求宽高结果有点尴尬NEON 版本的耗时只比标量版本低了大约 10%也就是从 0.42ms 降到 0.38ms 左右。原因也不难理解NMS 里到处都是分支判断if (suppressed[j])、if (a.class_id ! boxes[j].class_id)这些分支会打断 SIMD 的数据并行节奏。而且当候选框已经被削到 300 个左右时单帧 IoU 计算总量非常小NEON 的吞吐优势根本发挥不出来。我最后没有把 NEON 版本合入生产代码因为 0.04ms 的收益不值得后续维护成本。这里想借自己的经历提醒大家指令级优化要放在“数据规模确实很大且计算是主要瓶颈”时才做否则很容易变成自我感动。4.3 float 还是 double、内存对齐与 RKNN 输出的隐藏坑RK3588 的 A76 虽然有 FP64 硬件但吞吐通常只有 FP32 的一半A55 小核的 FP64 表现更差。NMS 坐标和 IoU 完全用float足够我实测在测试集上的结果与 Python double 版本几乎没有差异所以没必要用double。另外一个容易踩的坑是内存对齐。RKNN 工具包输出的 tensor 内存有时候是 int8 或 fp16 的量化数据不能拿过来直接当成float数组用。我建议在解码阶段把坐标先转成float再填充到自己定义的BBox结构体里。还有如果BBox结构体要反复访问可以考虑加alignas(16)或调整成员顺序让结构体在 16 字节边界对齐。RK3588 的 A76 对非对齐访问虽然没有 x86 那么敏感但某些访存场景下还是有性能惩罚的。我最后用的结构体里没有加alignas(16)因为测试结果差别不大。但如果你要做 NEON 优化对齐就是必须考虑的事否则vld1q_f32这类指令很可能直接报错。5. 正确性验证如何确保 C NMS 和 PyTorch 输出一致5.1 用 torchvision NMS 当基准误差容忍多少才合理优化后处理最怕的一件事速度上去了检测结果变了。所以我做了一套对比流程在开发机上用 PyTorch 的torchvision.ops.nms作为基准把同一批输入保存成二进制文件再喂给 C NMS对比最终保留框的 ID 集合。理论上两种实现完全一致时输出框的 ID 应该一模一样。但实际操作中会因为几个因素产生个别差异两个框的 IoU 恰好等于阈值时不同实现可能因为浮点误差一个判 true 一个判 false。分数相同的两个框排序稳定性会影响保留顺序。std::partial_sort和 PyTorch 内部排序策略不一定一致。所以我定的标准是单张图允许最多 3 个框 ID 差异且最终保留下来的框坐标必须严格一致。如果差异超过这个范围先查是不是 TopK 削掉了真值框再查 IoU 公式有没有写错。另外TopK 参数必须验证。我用验证集上 mAP 变化来确定 TopKK100 时掉了 0.8%K300 时基本不掉K500 和 K2000 几乎一样。所以最终选了 300既保证准确率又不给性能拖后腿。5.2 完整的加速数据表359 倍是怎么测出来的版本后处理策略平均候选框数单帧耗时加速倍数Python 循环版无预过滤直接双层循环 NMS25200152.4 ms1xC 朴素版sort erase 全量 NMS2520018.6 ms8.2xC TopK 优化置信度过滤 类别 TopK 300 partial_sort 标记数组平均 3502.85 ms53.5xC 最终版上述优化 4 线程 A76 并行平均 3500.42 ms359x这里的 359 倍是相对最开始的 Python 版本算的不是相对 C 朴素版。看数据能发现候选框削减和并行各自贡献了一大部分语言层面的优势反而是基础。我反复测了多次0.42ms 是中位数不是刻意挑出来的最好成绩。5.3 部署后持续踩的坑并发、内存分配和 topK 过小最终版本合入完整推理管线后又踩了几个新坑这里也一并记下来免得有人重复栽第一个坑是并发结果合并。最开始我做并行 NMS 时所有线程往同一个全局std::vectorBBox里写结果还加了一把 mutex。结果整个 NMS 比单线程还慢因为锁竞争和缓存行争抢太严重了。正确做法是每个线程维护自己的局部结果 vector最后再统一合并。第二个坑是内存分配。TopK 后候选框虽然少了但如果每个类别都频繁创建局部 vector还是会触发较多堆分配。我在处理前对 vector 做了reserve省掉不必要的扩容拷贝。第三个坑是 topK 过小导致漏检。如果你做的是密集场景比如行人聚集、货架上的密集商品每个类别 300 个候选框可能不够。我后来把这个值做成配置文件里的参数根据场景动态调而不是硬编码。第四个坑是大核与小核调度。OpenMP 不手动绑核时线程有时会被调度到 A55 小核上NMS 耗时波动明显。我用pthread_setaffinity_np把 4 个线程都绑到 CPU0-3A76上之后耗时才稳定下来。如果你的系统里还有其他后台任务抢大核还要再根据实际负载分配。第五个坑是散热降频。长稳测试时发现长时间跑 NPU 加 4 个 A76 满载RK3588 温度上来后局部频率会降NMS 从 0.42ms 可能变成 0.6ms 左右。虽然不至于崩但在严格实时场景里要注意留够余量。就我自己的经验来说NMS 从 150ms 降到 0.42ms最大的功臣不是 C 也不是编译器优化而是把候选框从 25200 个削减到 300 个这个“算法级降复杂度”的决定。如果当初我一头扎进 NEON 内联汇编估计也就能把 18ms 优化到 12ms 左右绝对到不了现在的量级。所以遇到后处理性能瓶颈先量数据规模、再看候选框冗余最后才谈指令级优化。这个顺序反了大概率会白忙一场。如果后续要做更高帧率或更大的检测模型我会把 TopK、线程绑定、甚至 NEON 这些手段再全部打开但每一步都会严格对着基准做回归测试。