1. 从 Python 后处理瓶颈说起为什么 NMS 值得单独拎出来优化在 RK3588 上跑 YOLOv5s模型推理本身其实没有想象中那么慢。RK3588 的 NPU 算力摆在那里INT8 量化后的 YOLOv5s 单帧推理通常能压到几十毫秒级别。但很多人部署完之后会发现一个很尴尬的现象推理快了整体帧率却上不去。问题往往不出在 NPU 上而是出在推理之后的那段后处理代码里。后处理里最耗时的环节十有八九是NMSNon-Maximum Suppression非极大值抑制。YOLOv5s 在三个尺度上输出预测框以 640x640 输入为例三个特征图分别是 80x80、40x40、20x20每个格子 3 个 anchor总共是 (80×80 40×40 20×20) × 3 25200 个候选框。经过置信度阈值过滤之后通常还剩几百到几千个框这些框都要参与 NMS 的 IoU 计算和排序。我最初用的是 Python 版本的 NMS逻辑清晰、写起来快但实测下来单帧后处理要 30ms 到 80ms 不等框多的时候甚至超过 100ms。这就导致一个很别扭的局面NPU 推理 40ms后处理 80ms整体帧率被后处理拖到了个位数。把 NMS 换成 C 实现之后同样的输入后处理时间直接掉到了 0.2ms 左右加速比达到 359 倍。这个数字不是理论值是我在同一块板子、同一组测试图片上反复跑出来的。这篇文章就把这条链路完整讲清楚Python NMS 到底慢在哪、C NMS 怎么写、怎么用 pybind11 把它接回 Python 流程、以及实测中那些容易翻车的细节。适合已经在 RK3588 上跑通了 YOLOv5s 推理、但被后处理卡住帧率的人也适合想理解 NMS 底层实现和 pybind11 混合编程的读者。2. Python NMS 到底慢在什么地方2.1 逐框循环 Python 解释器开销先看一段典型的 Python NMS 实现逻辑上没问题但性能上问题很大def nms_python(boxes, scores, iou_threshold): # boxes: [N, 4], scores: [N] order scores.argsort()[::-1] keep [] while order.size 0: i order[0] keep.append(i) if order.size 1: break # 计算当前框与剩余所有框的 IoU xx1 np.maximum(boxes[i, 0], boxes[order[1:], 0]) yy1 np.maximum(boxes[i, 1], boxes[order[1:], 1]) xx2 np.minimum(boxes[i, 2], boxes[order[1:], 2]) yy2 np.minimum(boxes[i, 3], boxes[order[1:], 3]) w np.maximum(0.0, xx2 - xx1) h np.maximum(0.0, yy2 - yy1) inter w * h area_i (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1]) area_rest (boxes[order[1:], 2] - boxes[order[1:], 0]) * \ (boxes[order[1:], 3] - boxes[order[1:], 1]) iou inter / (area_i area_rest - inter) inds np.where(iou iou_threshold)[0] order order[inds 1] return keep这段代码看起来用了 NumPy 向量化比纯 Python 循环快不少但它有两个致命问题。第一外层 while 循环是串行的每轮只处理一个框框多的时候循环次数就多。第二每轮循环都在做数组切片和重新分配order[1:]、np.where、order[inds 1]这些操作都会产生新的临时数组内存分配和拷贝的开销累积起来非常可观。2.2 内存分配与缓存不友好的代价Python 的 NumPy 数组底层虽然是连续的 C 内存但每次切片、np.maximum、np.where都会新建数组。假设有 3000 个候选框NMS 平均要循环几百轮每轮分配好几个临时数组总的内存分配次数就是几千次。在 RK3588 这种 ARM 平台上内存带宽和分配器性能都不如桌面 x86这些开销会被放大。更关键的是缓存局部性。Python 版本每轮只访问一个框boxes[i]然后和剩余所有框比较访问模式是跳跃的。而 C 版本可以把框按分数排好序之后用连续内存顺序访问CPU 缓存命中率高得多。这一点在候选框数量大的时候差异特别明显。2.3 GIL 与单线程限制Python 有全局解释器锁GIL纯 Python 代码的多线程没法真正并行。虽然后处理本身可以设计成多线程但在 Python 层面做这件事收益很低。C 实现天然没有这个限制后续如果想进一步并行化比如按类别分组并行 NMS扩展起来也方便。我实测过一组数据用同一张 640x640 的图片候选框数量约 2500 个置信度阈值 0.25IoU 阈值 0.45实现方式单帧后处理耗时相对加速比Python NumPy NMS约 72 ms1xC NMS单线程约 0.20 ms约 360xC NMS多线程4 核约 0.08 ms约 900x这个表格里的数字是多次测量取中位数的结果不同图片、不同阈值会有波动但量级差异是稳定的。359 倍这个数字就是来自 Python 72ms 对 C 0.2ms 的对比。3. C NMS 的实现思路与关键取舍3.1 数据结构选型为什么用 vector 而不是链表C NMS 的核心数据结构很简单一个std::vectorcv::Rect2f或者自定义的struct Box { float x1, y1, x2, y2, score; int label; }。用std::vector而不是链表原因是 NMS 需要频繁的随机访问和排序vector 的连续内存布局对缓存友好排序也快。实际实现里我推荐把框和分数打包成一个结构体排序的时候直接按 score 排避免维护两个平行数组struct Box { float x1, y1, x2, y2; float score; int label; }; // 按分数降序排序 std::sort(boxes.begin(), boxes.end(), [](const Box a, const Box b) { return a.score b.score; });这里有个细节排序是 NMS 里除了 IoU 计算之外的第二大开销。std::sort平均复杂度 O(N log N)2500 个框排序大概几十微秒可以接受。如果候选框数量特别大比如上万可以考虑用基数排序或者部分排序但一般 YOLOv5s 的场景用std::sort就够了。3.2 IoU 计算的向量化与提前终止IoU 计算是 NMS 的热点。每个保留框都要和剩余框算 IoU朴素实现是双重循环 O(N²)。优化手段有几个第一按分数排序后一旦遇到 IoU 超过阈值的框就把它从候选集里剔除这样候选集每轮都在缩小实际复杂度远低于 N²。第二用提前终止。如果当前框和某个候选框的 x 方向或 y 方向完全不重叠可以直接跳过面积计算。判断两个矩形是否相交只需要四次比较inline bool overlap(float x1, float y1, float x2, float y2, float x1_, float y1_, float x2_, float y2_) { return (x1 x2_ x2 x1_ y1 y2_ y2 y1_); }第三用 SIMD 指令批量计算。RK3588 是 ARM Cortex-A76/A55 架构支持 NEON 指令集。如果候选框数量很大可以用 NEON 一次算 4 个框的 IoU。不过实测下来对于 YOLOv5s 这种量级标量代码加上提前终止已经足够快NEON 优化带来的额外收益有限反而增加代码复杂度。我的建议是先把标量版本写对、写稳再考虑 SIMD。3.3 按类别分组 NMS 的处理YOLOv5s 是多类别检测NMS 必须按类别独立进行。如果所有类别的框混在一起做 NMS不同类别的框会互相抑制导致漏检。正确做法是先按 label 分组每组单独 NMS最后合并结果。分组的方式有两种一种是用std::unordered_mapint, std::vectorBox按 label 分桶另一种是先按 label 排序然后顺序处理每个类别段。前者代码直观后者缓存更友好。我实际用的是后者因为排序一次就能把同类别框聚在一起后续处理都是连续内存访问。std::sort(boxes.begin(), boxes.end(), [](const Box a, const Box b) { if (a.label ! b.label) return a.label b.label; return a.score b.score; }); std::vectorBox result; size_t start 0; while (start boxes.size()) { size_t end start; while (end boxes.size() boxes[end].label boxes[start].label) end; // 对 [start, end) 这一段做 NMS nms_single_class(boxes, start, end, iou_threshold, result); start end; }这个写法比 map 分桶少了一次内存分配实测在 2500 框规模下能省下十几微秒。4. pybind11 把 C NMS 接回 Python 流程4.1 为什么选 pybind11 而不是 ctypes 或 Cython把 C 函数暴露给 Python常见方案有 ctypes、Cython、SWIG、pybind11。我选 pybind11 的理由很直接ctypes需要手写 C 接口类型转换麻烦传 NumPy 数组要手动处理指针和 shape容易出错。Cython需要写 .pyx 文件编译流程多一层调试不如纯 C 直观。pybind11是 header-only 的直接 include 就能用对 NumPy 数组的支持通过pybind11/numpy.h开箱即用类型转换自动完成。对于 NMS 这种输入是 NumPy 数组、输出是索引列表的场景pybind11 的py::array_tfloat用起来非常顺手。4.2 接口设计输入输出怎么定接口设计要考虑两点一是 Python 侧调用方便二是 C 侧解析高效。我最终的接口是这样的#include pybind11/pybind11.h #include pybind11/numpy.h #include pybind11/stl.h namespace py pybind11; std::vectorint nms(py::array_tfloat boxes, py::array_tfloat scores, py::array_tint labels, float iou_threshold, float score_threshold) { auto boxes_buf boxes.request(); auto scores_buf scores.request(); auto labels_buf labels.request(); // ... 解析、过滤、排序、NMS ... return keep_indices; } PYBIND11_MODULE(nms_cpp, m) { m.doc() C NMS for YOLOv5s on RK3588; m.def(nms, nms, Non-Maximum Suppression, py::arg(boxes), py::arg(scores), py::arg(labels), py::arg(iou_threshold) 0.45f, py::arg(score_threshold) 0.25f); }输入用三个独立的 NumPy 数组boxes、scores、labels而不是一个打包的结构体数组。原因是 Python 侧从模型输出解析出来的就是分开的数组打包反而多一步。输出是std::vectorintpybind11 会自动转成 Python list。这里有个性能细节py::array_t默认会做类型检查和 shape 检查如果确定输入格式正确可以用py::array_tfloat, py::array::c_style | py::array::forcecast来减少检查开销。不过对于 NMS 这种调用频率每帧一次这点开销可以忽略。4.3 编译配置与 RK3588 上的注意事项编译用 CMake 最省事关键是找到 pybind11 和 Python 的头文件路径。在 RK3588 的 Ubuntu 系统上Python 通常是 3.8 或 3.10用python3-config --includes能拿到正确的 include 路径。cmake_minimum_required(VERSION 3.10) project(nms_cpp) set(CMAKE_CXX_STANDARD 14) set(CMAKE_CXX_STANDARD_REQUIRED ON) find_package(pybind11 REQUIRED) find_package(Python3 COMPONENTS Development REQUIRED) pybind11_add_module(nms_cpp nms.cpp) target_compile_options(nms_cpp PRIVATE -O3 -marchnative)-O3是必须的-marchnative让编译器针对 RK3588 的 Cortex-A76 生成优化指令。实测下来开-O3比-O2在 NMS 这种循环密集的代码上能快 20% 到 30%。注意如果编译时报Python.h not found检查python3-dev是否安装。RK3588 的板子有些精简系统默认不带开发头文件需要apt install python3-dev。编译完之后会生成nms_cpp.cpython-38-aarch64-linux-gnu.so这样的文件直接放在 Python 脚本同目录下import nms_cpp就能用。5. 实测数据与踩坑记录5.1 加速比是怎么测出来的测试环境RK3588 开发板Ubuntu 20.04Python 3.8YOLOv5s INT8 RKNN 模型输入 640x640。测试方法是用同一组 100 张图片分别跑 Python NMS 和 C NMS记录后处理阶段的耗时取平均值。测试代码的核心逻辑import time import nms_cpp # 假设 boxes, scores, labels 已经从模型输出解析好 t0 time.perf_counter() keep nms_cpp.nms(boxes, scores, labels, 0.45, 0.25) t1 time.perf_counter() print(fC NMS: {(t1 - t0) * 1000:.3f} ms)Python 版本用同样的输入同样的阈值。为了避免单次测量误差每张图跑 10 次取中位数。结果汇总指标Python NMSC NMS平均耗时72.3 ms0.201 ms最大耗时118.5 ms0.312 ms最小耗时41.2 ms0.156 ms加速比1x359xPython 版本耗时波动很大因为 NumPy 的临时数组分配受内存状态影响。C 版本非常稳定基本在 0.2ms 上下浮动。5.2 踩过的坑数组内存布局与 dtype第一个坑是NumPy 数组的内存布局。从 RKNN 模型输出解析出来的数组有时候是 C 连续row-major有时候是 Fortran 连续column-major甚至可能是非连续的切片。pybind11 的py::array_t默认要求 C 连续如果传进去的是非连续数组会抛异常或者静默拷贝。解决办法是在 Python 侧显式转换boxes np.ascontiguousarray(boxes, dtypenp.float32) scores np.ascontiguousarray(scores, dtypenp.float32) labels np.ascontiguousarray(labels, dtypenp.int32)np.ascontiguousarray保证内存连续dtype保证类型匹配。这一步看起来多余但不做的话在 RK3588 上会随机崩溃而且崩溃位置不固定非常难查。第二个坑是dtype 不匹配。RKNN 输出的分数有时候是 float16有时候是 float32如果 C 侧声明的是py::array_tfloat而实际传的是 float16pybind11 会尝试转换转换失败就报错。我的做法是在 Python 侧统一转成 float32C 侧只接受 float32。5.3 踩过的坑阈值边界与空输入第三个坑是空输入。如果某一帧没有任何框超过置信度阈值boxes数组的 shape 是(0, 4)这时候 C 侧的排序和循环要能正确处理空数组。我最初的实现里boxes[0]直接越界导致段错误。修复方法是在函数开头加判断if (boxes_buf.shape[0] 0) return {};第四个坑是IoU 阈值的边界条件。当两个框完全重合时IoU 是 1.0如果阈值设成 1.0iou threshold为 false两个框都会保留。实际使用中阈值一般设 0.45 到 0.5不会碰到这个边界。但如果做单元测试要覆盖 IoU 等于阈值的情况判断条件用还是会影响结果。我统一用即 IoU 严格大于阈值才抑制。5.4 踩过的坑多线程与线程安全如果后续想把 NMS 放到多线程里跑比如多个摄像头流并行要注意 C 函数本身是线程安全的因为它不依赖任何全局状态所有数据都在栈上或参数里。但 pybind11 的模块初始化是单次的多线程调用同一个函数没问题。不过有个细节Python 的 GIL 在调用 C 函数时默认是持有的。如果 C 函数执行时间很长会阻塞其他 Python 线程。对于 NMS 这种 0.2ms 的函数无所谓。但如果以后扩展成更复杂的后处理可以在 pybind11 绑定里加py::call_guardpy::gil_scoped_release()来释放 GILm.def(nms, nms, py::call_guardpy::gil_scoped_release());加了之后 C 执行期间不持有 GIL其他 Python 线程可以继续跑。但要注意释放 GIL 后不能再操作 Python 对象所以这个函数里不能碰py::array_t之外的东西。6. 从 NMS 延伸出去的优化空间6.1 把后处理整体搬到 CNMS 只是后处理的一环。完整的 YOLOv5s 后处理还包括从模型输出解码边界框anchor 解码、sigmoid 激活、置信度过滤、坐标缩放。这些步骤如果都在 Python 里做累积开销也不小。我后来的做法是把整个后处理链路都写成 CPython 只负责调用和拿结果。这样做的好处是减少 Python 和 C 之间的数据往返。如果只把 NMS 放 C那解码和过滤还在 Python中间还是要传大数组。整体搬过去之后Python 侧只需要传原始模型输出C 侧一次性完成所有处理返回最终框。6.2 与 RKNN 推理流程的衔接RKNN 的 Python API 返回的输出是 NumPy 数组直接喂给 C 后处理函数即可。如果用的是 RKNN 的 C API那整个流程都在 C 里连 pybind11 都不需要。但对于快速原型开发Python pybind11 的组合更灵活调试也方便。实际部署时如果追求极致性能可以把推理和后处理都放在 C 里Python 只做业务逻辑。如果追求开发效率Python 调 C 后处理已经能拿到 359 倍的加速性价比很高。6.3 进一步压榨NEON 与多线程前面提到 NEON 优化收益有限但在候选框数量特别大比如 10000 以上的场景下NEON 批量计算 IoU 还是有意义的。另外如果一帧里有多个类别可以按类别分组后用多线程并行 NMSRK3588 有 8 个核理论上能再快几倍。不过要注意多线程的收益受限于候选框数量和核间通信开销。2500 个框分到 4 个线程每个线程处理 600 多个框线程创建和同步的开销可能吃掉一部分收益。实测 4 线程比单线程快 2 到 3 倍没有达到理想的 4 倍就是因为开销。我的建议是先用单线程 C 版本确认瓶颈确实在 NMS 上再考虑多线程。很多时候单线程 C 已经把后处理压到 0.2ms相对于推理的 40ms 已经可以忽略没必要过度优化。7. 一些实际部署中的经验在 RK3588 上做这套东西有几个经验值得单独说。第一编译产物要跟 Python 版本匹配。RK3588 的板子上可能同时装了 Python 3.8 和 3.10pybind11 编译出来的 .so 文件名里带版本号import的时候要确保用的是对应版本的解释器。我踩过一次坑用 3.8 编译的模块在 3.10 里 import 报错查了半天才发现是版本问题。第二交叉编译 vs 板端编译。如果是在 x86 主机上交叉编译要确保 CMake 里的 Python 路径指向目标板的 Python而不是主机的。我一般直接在板子上编译虽然慢一点但省去了交叉编译工具链配置的麻烦。RK3588 的板子编译一个几百行的 C 文件也就十几秒。第三内存对齐。RK3588 的 NPU 输出有时候是 16 字节对齐的如果 C 侧按普通数组处理可能触发未对齐访问。虽然 ARM 一般能处理未对齐访问但性能会下降。用py::array_t的时候pybind11 会保证对齐不用额外操心。第四测试要覆盖边界情况。空输入、单框、所有框 IoU 都超过阈值、所有框 IoU 都不超过阈值这些情况都要测。我写了一个小的测试脚本构造各种极端输入确保 C 函数不崩溃、结果正确。import numpy as np import nms_cpp # 空输入 assert nms_cpp.nms(np.zeros((0, 4), np.float32), np.zeros((0,), np.float32), np.zeros((0,), np.int32), 0.45, 0.25) [] # 单框 boxes np.array([[0, 0, 10, 10]], np.float32) scores np.array([0.9], np.float32) labels np.array([0], np.int32) assert nms_cpp.nms(boxes, scores, labels, 0.45, 0.25) [0] # 完全重合的两个框 boxes np.array([[0, 0, 10, 10], [0, 0, 10, 10]], np.float32) scores np.array([0.9, 0.8], np.float32) labels np.array([0, 0], np.int32) assert nms_cpp.nms(boxes, scores, labels, 0.45, 0.25) [0]这几个测试跑通基本能保证函数在各种输入下行为正确。最后说一个我自己的体会性能优化要先测量再动手。我一开始以为瓶颈在 NPU 推理花了很多时间调模型量化结果发现后处理才是大头。用time.perf_counter()把推理和后处理分开计时一眼就能看出问题在哪。C NMS 这个优化从动手到跑通大概花了一个下午但带来的帧率提升是立竿见影的。如果你也在 RK3588 上跑 YOLOv5s并且被帧率卡住先看看后处理耗时很可能就是 NMS 在拖后腿。
