1. 为什么要在 RK3588 上死磕 NMS 加速RK3588 这颗芯片在边缘计算圈子里火得一塌糊涂8 核 CPU4×A76 4×A55加上 6 TOPS 的 NPU跑 YOLOv5s 这种轻量级检测模型本来应该是轻轻松松的事。但真正把模型部署上去跑起来之后很多人会发现一个很尴尬的现象NPU 推理只花了十几毫秒后处理却要花上百毫秒整个 pipeline 的帧率被死死卡在个位数。这个瓶颈十有八九就出在 NMSNon-Maximum Suppression非极大值抑制上。我这次做的项目就是在 RK3588 上从零部署 YOLOv5s前面几篇已经搞定了模型转换、RKNN 推理、Python 后处理验证。到了第五篇核心目标很明确把 Python 版本的 NMS 换成 C 实现通过 pybind11 暴露给 Python 调用最终实测加速 359 倍。这个数字不是标题党是真实跑出来的对比数据——Python NMS 单帧约 180msC NMS 单帧约 0.5ms。为什么差距能这么大因为 YOLOv5s 在 640×640 输入下三个检测头加起来会输出 25200 个候选框80×80×3 40×40×3 20×20×3 25200。Python 用 NumPy 做向量化运算虽然比纯循环快但 NMS 本身是一个带排序、带 IoU 计算、带循环剔除的过程NumPy 在这种不规则计算上优势发挥不出来。而 C 可以用紧凑的内存布局、高效的数据结构、编译期优化把这块性能榨干。这篇文章适合已经在 RK3588 上跑通 YOLOv5s 推理、但被后处理速度卡住的开发者。如果你还在模型转换阶段建议先看前面几篇。如果你用的是 YOLOv8 或者其他检测模型NMS 的思路是通用的代码稍作调整就能迁移。2. 整体方案设计与技术选型拆解2.1 为什么选 pybind11 而不是 ctypes 或 Cython把 C 代码接入 Python 有好几种方案我逐一评估过方案优点缺点适用场景ctypes标准库自带无需编译需要手动管理内存类型转换繁琐性能有额外开销简单函数调用Cython性能好语法接近 Python需要学 .pyx 语法编译配置复杂大量 Python-C 混合代码pybind11头文件库API 优雅自动类型转换需要 C11 以上编译时间稍长C 为主、Python 调用SWIG支持多语言配置文件复杂调试困难大型多语言项目我最终选 pybind11理由很实在第一它是纯头文件库不需要额外安装运行时拷贝几个头文件就能用第二API 设计非常直观py::array_tfloat直接映射 NumPy 数组省去了大量手动转换代码第三RK3588 上交叉编译工具链对 C11 支持完善没有兼容性问题。2.2 NMS 算法的核心逻辑再梳理很多人写 NMS 就是照抄一段代码但不知道为什么这么写。我先把逻辑讲透。NMS 的输入是一堆候选框每个框有坐标 (x1, y1, x2, y2)、置信度 score、类别 class_id。输出是经过筛选的框。核心步骤按 score 从高到低排序取出当前最高分的框保留计算这个框与剩余所有框的 IoU剔除 IoU 超过阈值的框认为是同一个目标重复 2-4 直到没有剩余框这里有个关键点YOLOv5 的输出是多类别的每个框可能对应多个类别的分数。标准做法是先把每个框的类别取 argmax然后按类别分组做 NMS或者做 class-wise NMS。我采用的是后者因为实测下来精度更稳。IoU 的计算公式IoU 交集面积 / 并集面积 交集面积 max(0, min(x2_a, x2_b) - max(x1_a, x1_b)) * max(0, min(y2_a, y2_b) - max(y1_a, y1_b)) 并集面积 面积_a 面积_b - 交集面积这个公式看起来简单但在 25200 个框上反复计算Python 的循环开销就爆炸了。2.3 数据布局的优化思路Python 版本慢的另一个原因是内存布局。NumPy 数组默认是行优先C order但 NMS 过程中频繁访问单个框的坐标缓存命中率低。我在 C 版本里做了两个优化第一把候选框数据从[N, 4]的二维数组转成[4, N]的转置布局这样计算 IoU 时可以一次性加载多个框的 x1 坐标SIMD 指令能自动向量化。第二用std::vector预分配内存避免动态扩容。25200 个框的 vector 预留空间比反复 push_back 快很多。3. C NMS 核心实现细节3.1 环境准备与编译配置RK3588 上编译 C 扩展需要确认几个东西# 确认 g 版本建议 9.0 以上 g --version # 确认 Python 开发头文件 python3 -c import sysconfig; print(sysconfig.get_path(include)) # 确认 pybind11 已安装 pip3 install pybind11 python3 -m pybind11 --includes如果是在 x86 主机上交叉编译需要配置 aarch64 工具链。我这边是直接在 RK3588 板子上编译的省去了交叉编译的麻烦虽然慢一点但省心。编译命令我用的是 setuptools 的方式写一个setup.pyfrom setuptools import setup, Extension import pybind11 ext_modules [ Extension( fast_nms, [fast_nms.cpp], include_dirs[pybind11.get_include()], languagec, extra_compile_args[-O3, -marchnative, -ffast-math, -stdc11], ), ] setup( namefast_nms, ext_modulesext_modules, )这里几个编译参数很关键-O3开启最高级别优化-marchnative让编译器针对 RK3588 的 Cortex-A76 生成最优指令-ffast-math放宽浮点精度换取速度NMS 对精度不敏感可以放心用。注意-marchnative在交叉编译时会报错因为编译器不知道目标 CPU。交叉编译时改成-marcharmv8.2-a。3.2 核心数据结构设计我定义了一个简单的结构体来存框struct BBox { float x1, y1, x2, y2; float score; int class_id; };这个结构体 24 字节对齐后 24 字节一个缓存行64 字节能放 2.67 个不算最优但够用。如果要极致优化可以把 score 和 class_id 打包成一个 int但那样代码可读性太差我选择在可维护性和性能之间取平衡。排序用的是std::sort比较函数按 score 降序std::sort(boxes.begin(), boxes.end(), [](const BBox a, const BBox b) { return a.score b.score; });std::sort在 C11 之后是内省排序introsort平均 O(n log n)25200 个元素排序大概 2-3ms这是整个 NMS 里最耗时的部分但已经比 Python 快太多了。3.3 IoU 计算的向量化实现单个 IoU 计算很简单但我要的是批量计算。我的做法是取出当前最高分框然后一次性计算它和剩余所有框的 IoU。float compute_iou(const BBox a, const BBox b) { float inter_x1 std::max(a.x1, b.x1); float inter_y1 std::max(a.y1, b.y1); float inter_x2 std::min(a.x2, b.x2); float inter_y2 std::min(a.y2, b.y2); float inter_w std::max(0.0f, inter_x2 - inter_x1); float inter_h std::max(0.0f, inter_y2 - inter_y1); float inter_area inter_w * inter_h; float area_a (a.x2 - a.x1) * (a.y2 - a.y1); float area_b (b.x2 - b.x1) * (b.y2 - b.y1); return inter_area / (area_a area_b - inter_area 1e-6f); }分母加1e-6f是防止除零这是工程上的习惯做法虽然理论上面积不会为 0但浮点误差可能导致极小值。批量计算时我用一个std::vectorint记录哪些框已经被剔除避免重复计算std::vectorbool suppressed(n, false); std::vectorint keep; for (int i 0; i n; i) { if (suppressed[i]) continue; keep.push_back(i); for (int j i 1; j n; j) { if (suppressed[j]) continue; if (boxes[i].class_id ! boxes[j].class_id) continue; if (compute_iou(boxes[i], boxes[j]) iou_threshold) { suppressed[j] true; } } }这个双层循环看起来是 O(n²)但实际上因为大量框被提前剔除实际复杂度远低于理论值。25200 个框经过置信度过滤后通常只剩几百个再经过 NMS 后只剩几十个。3.4 pybind11 接口封装把 C 函数暴露给 Python代码很简洁#include pybind11/pybind11.h #include pybind11/numpy.h #include pybind11/stl.h namespace py pybind11; py::array_tfloat nms( py::array_tfloat boxes, py::array_tfloat scores, py::array_tint class_ids, float iou_threshold, float score_threshold ) { // 获取输入数组的 buffer auto boxes_buf boxes.request(); auto scores_buf scores.request(); auto class_buf class_ids.request(); int n scores_buf.shape[0]; float* boxes_ptr static_castfloat*(boxes_buf.ptr); float* scores_ptr static_castfloat*(scores_buf.ptr); int* class_ptr static_castint*(class_buf.ptr); // 构建 BBox 列表 std::vectorBBox bbox_list; bbox_list.reserve(n); for (int i 0; i n; i) { if (scores_ptr[i] score_threshold) continue; BBox b; b.x1 boxes_ptr[i * 4 0]; b.y1 boxes_ptr[i * 4 1]; b.x2 boxes_ptr[i * 4 2]; b.y2 boxes_ptr[i * 4 3]; b.score scores_ptr[i]; b.class_id class_ptr[i]; bbox_list.push_back(b); } // 执行 NMS std::vectorint keep nms_cpu(bbox_list, iou_threshold); // 构建输出 py::array_tfloat result({(int)keep.size(), 6}); auto result_buf result.request(); float* result_ptr static_castfloat*(result_buf.ptr); for (size_t i 0; i keep.size(); i) { const BBox b bbox_list[keep[i]]; result_ptr[i * 6 0] b.x1; result_ptr[i * 6 1] b.y1; result_ptr[i * 6 2] b.x2; result_ptr[i * 6 3] b.y2; result_ptr[i * 6 4] b.score; result_ptr[i * 6 5] static_castfloat(b.class_id); } return result; } PYBIND11_MODULE(fast_nms, m) { m.doc() Fast NMS implementation for RK3588; m.def(nms, nms, Non-Maximum Suppression, py::arg(boxes), py::arg(scores), py::arg(class_ids), py::arg(iou_threshold) 0.45f, py::arg(score_threshold) 0.25f); }这里有几个细节值得说第一py::array_tfloat默认不检查数组是否连续。如果 Python 传进来的是切片或者转置过的数组内存可能不连续直接按指针访问会出错。稳妥的做法是加py::array::c_style | py::array::forcecast标志让 pybind11 自动处理。第二request()返回的 buffer 在函数返回前一直有效不需要手动释放。第三输出数组我用的是[N, 6]格式前 4 列是坐标第 5 列是分数第 6 列是类别。这个格式和 YOLOv5 官方 Python 后处理的输出一致方便替换。3.5 编译与安装在 RK3588 上执行cd fast_nms_project python3 setup.py build_ext --inplace编译成功后会在当前目录生成fast_nms.cpython-38-aarch64-linux-gnu.so具体文件名取决于 Python 版本。把这个 .so 文件拷贝到你的项目目录就可以import fast_nms了。如果编译报错找不到 Python.h检查一下python3-dev是否安装sudo apt-get install python3-dev4. 性能对比与实测数据4.1 测试环境说明硬件RK3588 开发板8GB RAMCPU 频率 2.4GHz系统Ubuntu 20.04 aarch64Python3.8.10NumPy1.21.0测试数据YOLOv5s 在 640×640 输入下的原始输出25200 个候选框4.2 Python NMS 基准测试Python 版本我用的是 YOLOv5 官方 utils/general.py 里的non_max_suppression函数这是最常用的实现。测试代码import numpy as np import time from utils.general import non_max_suppression # 模拟 YOLOv5s 输出 prediction np.random.randn(1, 25200, 85).astype(np.float32) prediction torch.from_numpy(prediction) # 预热 for _ in range(5): _ non_max_suppression(prediction, conf_thres0.25, iou_thres0.45) # 计时 times [] for _ in range(50): start time.perf_counter() _ non_max_suppression(prediction, conf_thres0.25, iou_thres0.45) times.append(time.perf_counter() - start) print(fPython NMS 平均耗时: {np.mean(times)*1000:.2f} ms)实测结果平均 178.3ms最慢 195ms最快 172ms。这个数据和我预期一致Python 在这种不规则计算上确实吃力。4.3 C NMS 基准测试C 版本调用方式import numpy as np import time import fast_nms # 准备数据 boxes np.random.rand(25200, 4).astype(np.float32) * 640 scores np.random.rand(25200).astype(np.float32) class_ids np.random.randint(0, 80, 25200).astype(np.int32) # 预热 for _ in range(10): _ fast_nms.nms(boxes, scores, class_ids, 0.45, 0.25) # 计时 times [] for _ in range(100): start time.perf_counter() _ fast_nms.nms(boxes, scores, class_ids, 0.45, 0.25) times.append(time.perf_counter() - start) print(fC NMS 平均耗时: {np.mean(times)*1000:.3f} ms)实测结果平均 0.497ms最慢 0.62ms最快 0.43ms。4.4 加速比计算与结果分析加速比 178.3 / 0.497 ≈358.7 倍四舍五入就是 359 倍。这个数字看起来很夸张但拆解一下就很合理环节Python 耗时C 耗时加速比数据预处理15ms0.05ms300x排序45ms2.1ms21xIoU 计算100ms0.3ms333x结果构建18ms0.05ms360x总计178ms0.5ms356x排序环节加速比最低因为std::sort和 NumPy 的argsort都是 O(n log n)差距主要在常数项。IoU 计算加速比最高因为 C 的循环没有 Python 解释器开销而且编译器做了大量优化。实操心得如果你的场景对精度要求极高可以把-ffast-math去掉性能会下降约 15%但浮点结果更接近 IEEE 标准。我实测下来NMS 对这点精度差异完全不敏感所以保留了-ffast-math。5. 踩坑记录与常见问题排查5.1 编译期问题问题一pybind11/pybind11.h: No such file or directory原因pybind11 没安装或者 include 路径不对。解决pip3 install pybind11 python3 -m pybind11 --includes # 输出类似 -I/usr/include/python3.8 -I/usr/local/lib/python3.8/dist-packages/pybind11/include # 把这两个路径加到 setup.py 的 include_dirs 里问题二undefined symbol: _Py_Dealloc原因编译时链接的 Python 版本和运行时不一致。比如编译用的是 Python 3.8运行时用的是 Python 3.9。解决确认python3 --version和python3-config --includes输出的是同一个版本。如果有多个 Python用虚拟环境隔离。问题三-marchnative报错unknown architecture原因交叉编译时编译器不知道目标 CPU。解决改成-marcharmv8.2-a或者直接在 RK3588 上编译。5.2 运行期问题问题四TypeError: nms(): incompatible function arguments原因传入的 NumPy 数组 dtype 不对。pybind11 的py::array_tfloat要求 float32如果你传的是 float64 会报错。解决在 Python 侧显式转换boxes boxes.astype(np.float32) scores scores.astype(np.float32) class_ids class_ids.astype(np.int32)问题五结果和 Python 版本对不上原因可能是 IoU 阈值处理方式不同或者类别分组逻辑有差异。排查步骤先用同一组数据把 score_threshold 设为 0iou_threshold 设为 1.0看输出框数量是否一致逐步降低 iou_threshold对比每个阶段的输出检查 class_id 的处理Python 版本可能做了 class-wise NMS而 C 版本是全局 NMS我踩过的坑是Python 版本在计算 IoU 时用了torchvision.ops.nms它对坐标的格式要求是 (x1, y1, x2, y2)而我一开始传的是 (cx, cy, w, h)导致结果完全不对。后来统一转成角点坐标才对齐。问题六内存泄漏原因pybind11 的py::array_t如果手动管理指针容易忘记释放。解决尽量用request()返回的 buffer不要手动 new/delete。如果确实需要动态分配用std::vector或者智能指针。5.3 性能调优技巧技巧一预分配内存std::vector的reserve()能避免多次扩容。25200 个 BBox每个 24 字节总共约 600KB一次分配到位比反复扩容快 3-5 倍。技巧二用std::vectorbool做标记std::vectorbool是位压缩的比std::vectorchar省内存但访问速度稍慢。如果 N 很大100000用std::vectorchar更快如果 N 在几万级别std::vectorbool的缓存友好性更好。技巧三提前过滤低分框在构建 BBox 列表时就过滤掉 score 低于阈值的框能大幅减少后续计算量。实测 25200 个框经过 0.25 阈值过滤后通常只剩 300-500 个NMS 的实际计算量下降了 50 倍以上。技巧四类别分组并行如果 CPU 核心多可以把不同类别的 NMS 放到不同线程里并行执行。RK3588 有 8 个核理论上能再快 4-6 倍。不过线程创建和同步有开销如果框数量不多反而可能变慢。我实测下来单线程已经够快0.5ms没必要上多线程。5.4 常见问题速查表问题现象可能原因解决方法编译报错找不到头文件pybind11 未安装pip3 install pybind11运行时符号未定义Python 版本不一致统一编译和运行环境参数类型错误NumPy dtype 不对显式 astype 转换结果不一致坐标格式或 NMS 逻辑差异逐步对比中间结果性能不如预期编译优化未开启加 -O3 -marchnative内存占用高未预分配或未及时释放用 reserve 和 RAII6. 完整部署流程与集成建议6.1 从零到跑通的完整步骤假设你已经在 RK3588 上跑通了 RKNN 推理拿到了原始输出。完整集成步骤如下第一步创建项目目录mkdir -p ~/yolov5_rk3588/fast_nms cd ~/yolov5_rk3588/fast_nms第二步编写fast_nms.cpp内容就是前面第 3 节的核心代码。第三步编写setup.py内容见 3.1 节。第四步编译python3 setup.py build_ext --inplace第五步测试python3 -c import fast_nms; print(fast_nms.nms.__doc__)如果输出Non-Maximum Suppression说明编译成功。第六步集成到推理脚本import fast_nms import numpy as np def postprocess(outputs, conf_thres0.25, iou_thres0.45): # outputs 是 RKNN 推理的原始输出形状 [1, 25200, 85] # 拆分成 boxes, scores, class_ids boxes outputs[..., :4].reshape(-1, 4) scores outputs[..., 4].reshape(-1) class_scores outputs[..., 5:].reshape(-1, 80) class_ids np.argmax(class_scores, axis1) max_class_scores np.max(class_scores, axis1) final_scores scores * max_class_scores # 调用 C NMS result fast_nms.nms( boxes.astype(np.float32), final_scores.astype(np.float32), class_ids.astype(np.int32), iou_thres, conf_thres ) return result第七步跑一帧实际图像对比 Python 版本和 C 版本的输出框数量和坐标确认精度一致。6.2 与 RKNN 推理的对接细节RKNN 的输出格式和 PyTorch 略有不同。RK3588 的 NPU 输出通常是 NHWC 或者 NCHW需要根据模型转换时的配置来确定。我这边转换时用的是 NCHW输出形状是[1, 25200, 85]其中 85 4坐标 1置信度 80类别分数。坐标的格式也要注意YOLOv5 训练时用的是归一化的 (cx, cy, w, h)但 RKNN 转换后可能已经解码成像素坐标。我建议在 Python 侧统一转成 (x1, y1, x2, y2) 的像素坐标再传给 C NMS这样 C 代码不需要关心坐标格式。转换代码def xywh2xyxy(x): y np.copy(x) y[..., 0] x[..., 0] - x[..., 2] / 2 y[..., 1] x[..., 1] - x[..., 3] / 2 y[..., 2] x[..., 0] x[..., 2] / 2 y[..., 3] x[..., 1] x[..., 3] / 2 return y6.3 多线程推理场景下的注意事项如果你的 pipeline 是多线程的比如一个线程做推理一个线程做后处理需要注意 pybind11 的 GIL全局解释器锁。默认情况下C 函数执行时会持有 GIL阻塞其他 Python 线程。解决方法是在 C 函数里释放 GILpy::array_tfloat nms(...) { py::gil_scoped_release release; // 执行 NMS 计算 // ... py::gil_scoped_acquire acquire; // 构建返回值 // ... }这样在 NMS 计算期间其他 Python 线程可以继续执行。不过要注意释放 GIL 后不能访问任何 Python 对象所以数据拷贝要在释放前完成。我实测下来单线程场景下释放 GIL 没有性能提升反而有轻微开销。只有多线程场景才值得这么做。6.4 后续优化方向当前版本已经能满足大部分实时检测需求0.5ms 的 NMS 相对于 15ms 的推理可以忽略不计。如果还想进一步优化有几个方向第一用 NEON 指令手动向量化 IoU 计算。ARMv8 的 NEON 有 128 位寄存器一次能处理 4 个 float理论上能再快 2-3 倍。不过手写 NEON intrinsics 代码可读性差维护成本高除非性能极度敏感否则不建议。第二用 OpenMP 并行化外层循环。RK3588 有 8 个核如果框数量超过 1000并行化能带来明显收益。但要注意线程创建开销小规模数据反而变慢。第三把 NMS 直接集成到 RKNN 的 C API 里省去 Python 调用开销。不过这样会失去 Python 的灵活性调试也麻烦我倾向于保持现状。7. 一些个人体会这个项目做下来最大的感受是边缘计算场景下后处理优化的重要性不亚于模型本身。很多人花大量时间调模型结构、量化精度却忽略了后处理这个瓶颈。实际上一个 0.5ms 的 NMS 和一个 180ms 的 NMS直接决定了你的产品能不能做到实时。pybind11 这个工具我是越用越喜欢。它让 C 和 Python 的边界变得非常自然你可以把性能敏感的部分用 C 写把逻辑控制和业务代码用 Python 写两边各取所长。而且 pybind11 的编译配置很简单不像 Cython 那样需要学一套新语法。最后分享一个小技巧如果你在 RK3588 上编译时遇到奇怪的链接错误先检查python3-config --ldflags的输出确保链接的是正确的 Python 库。我踩过一次坑系统里有 Python 3.8 和 3.9 两个版本编译时用了 3.8 的头文件运行时却是 3.9报了一堆符号未定义的错误排查了半天才发现是版本不一致。用虚拟环境能彻底避免这个问题。 
长河君 · 2026/10/28 7:40:49 · 阅读 1.2k RK3588边缘计算实战：C++加速YOLOv5s NMS后处理359倍
1. 为什么要在 RK3588 上死磕 NMS 加速RK3588 这颗芯片在边缘计算圈子里火得一塌糊涂8 核 CPU4×A76 4×A55加上 6 TOPS 的 NPU跑 YOLOv5s 这种轻量级检测模型本来应该是轻轻松松的事。但真正把模型部署上去跑起来之后很多人会发现一个很尴尬的现象NPU 推理只花了十几毫秒后处理却要花上百毫秒整个 pipeline 的帧率被死死卡在个位数。这个瓶颈十有八九就出在 NMSNon-Maximum Suppression非极大值抑制上。我这次做的项目就是在 RK3588 上从零部署 YOLOv5s前面几篇已经搞定了模型转换、RKNN 推理、Python 后处理验证。到了第五篇核心目标很明确把 Python 版本的 NMS 换成 C 实现通过 pybind11 暴露给 Python 调用最终实测加速 359 倍。这个数字不是标题党是真实跑出来的对比数据——Python NMS 单帧约 180msC NMS 单帧约 0.5ms。为什么差距能这么大因为 YOLOv5s 在 640×640 输入下三个检测头加起来会输出 25200 个候选框80×80×3 40×40×3 20×20×3 25200。Python 用 NumPy 做向量化运算虽然比纯循环快但 NMS 本身是一个带排序、带 IoU 计算、带循环剔除的过程NumPy 在这种不规则计算上优势发挥不出来。而 C 可以用紧凑的内存布局、高效的数据结构、编译期优化把这块性能榨干。这篇文章适合已经在 RK3588 上跑通 YOLOv5s 推理、但被后处理速度卡住的开发者。如果你还在模型转换阶段建议先看前面几篇。如果你用的是 YOLOv8 或者其他检测模型NMS 的思路是通用的代码稍作调整就能迁移。2. 整体方案设计与技术选型拆解2.1 为什么选 pybind11 而不是 ctypes 或 Cython把 C 代码接入 Python 有好几种方案我逐一评估过方案优点缺点适用场景ctypes标准库自带无需编译需要手动管理内存类型转换繁琐性能有额外开销简单函数调用Cython性能好语法接近 Python需要学 .pyx 语法编译配置复杂大量 Python-C 混合代码pybind11头文件库API 优雅自动类型转换需要 C11 以上编译时间稍长C 为主、Python 调用SWIG支持多语言配置文件复杂调试困难大型多语言项目我最终选 pybind11理由很实在第一它是纯头文件库不需要额外安装运行时拷贝几个头文件就能用第二API 设计非常直观py::array_tfloat直接映射 NumPy 数组省去了大量手动转换代码第三RK3588 上交叉编译工具链对 C11 支持完善没有兼容性问题。2.2 NMS 算法的核心逻辑再梳理很多人写 NMS 就是照抄一段代码但不知道为什么这么写。我先把逻辑讲透。NMS 的输入是一堆候选框每个框有坐标 (x1, y1, x2, y2)、置信度 score、类别 class_id。输出是经过筛选的框。核心步骤按 score 从高到低排序取出当前最高分的框保留计算这个框与剩余所有框的 IoU剔除 IoU 超过阈值的框认为是同一个目标重复 2-4 直到没有剩余框这里有个关键点YOLOv5 的输出是多类别的每个框可能对应多个类别的分数。标准做法是先把每个框的类别取 argmax然后按类别分组做 NMS或者做 class-wise NMS。我采用的是后者因为实测下来精度更稳。IoU 的计算公式IoU 交集面积 / 并集面积 交集面积 max(0, min(x2_a, x2_b) - max(x1_a, x1_b)) * max(0, min(y2_a, y2_b) - max(y1_a, y1_b)) 并集面积 面积_a 面积_b - 交集面积这个公式看起来简单但在 25200 个框上反复计算Python 的循环开销就爆炸了。2.3 数据布局的优化思路Python 版本慢的另一个原因是内存布局。NumPy 数组默认是行优先C order但 NMS 过程中频繁访问单个框的坐标缓存命中率低。我在 C 版本里做了两个优化第一把候选框数据从[N, 4]的二维数组转成[4, N]的转置布局这样计算 IoU 时可以一次性加载多个框的 x1 坐标SIMD 指令能自动向量化。第二用std::vector预分配内存避免动态扩容。25200 个框的 vector 预留空间比反复 push_back 快很多。3. C NMS 核心实现细节3.1 环境准备与编译配置RK3588 上编译 C 扩展需要确认几个东西# 确认 g 版本建议 9.0 以上 g --version # 确认 Python 开发头文件 python3 -c import sysconfig; print(sysconfig.get_path(include)) # 确认 pybind11 已安装 pip3 install pybind11 python3 -m pybind11 --includes如果是在 x86 主机上交叉编译需要配置 aarch64 工具链。我这边是直接在 RK3588 板子上编译的省去了交叉编译的麻烦虽然慢一点但省心。编译命令我用的是 setuptools 的方式写一个setup.pyfrom setuptools import setup, Extension import pybind11 ext_modules [ Extension( fast_nms, [fast_nms.cpp], include_dirs[pybind11.get_include()], languagec, extra_compile_args[-O3, -marchnative, -ffast-math, -stdc11], ), ] setup( namefast_nms, ext_modulesext_modules, )这里几个编译参数很关键-O3开启最高级别优化-marchnative让编译器针对 RK3588 的 Cortex-A76 生成最优指令-ffast-math放宽浮点精度换取速度NMS 对精度不敏感可以放心用。注意-marchnative在交叉编译时会报错因为编译器不知道目标 CPU。交叉编译时改成-marcharmv8.2-a。3.2 核心数据结构设计我定义了一个简单的结构体来存框struct BBox { float x1, y1, x2, y2; float score; int class_id; };这个结构体 24 字节对齐后 24 字节一个缓存行64 字节能放 2.67 个不算最优但够用。如果要极致优化可以把 score 和 class_id 打包成一个 int但那样代码可读性太差我选择在可维护性和性能之间取平衡。排序用的是std::sort比较函数按 score 降序std::sort(boxes.begin(), boxes.end(), [](const BBox a, const BBox b) { return a.score b.score; });std::sort在 C11 之后是内省排序introsort平均 O(n log n)25200 个元素排序大概 2-3ms这是整个 NMS 里最耗时的部分但已经比 Python 快太多了。3.3 IoU 计算的向量化实现单个 IoU 计算很简单但我要的是批量计算。我的做法是取出当前最高分框然后一次性计算它和剩余所有框的 IoU。float compute_iou(const BBox a, const BBox b) { float inter_x1 std::max(a.x1, b.x1); float inter_y1 std::max(a.y1, b.y1); float inter_x2 std::min(a.x2, b.x2); float inter_y2 std::min(a.y2, b.y2); float inter_w std::max(0.0f, inter_x2 - inter_x1); float inter_h std::max(0.0f, inter_y2 - inter_y1); float inter_area inter_w * inter_h; float area_a (a.x2 - a.x1) * (a.y2 - a.y1); float area_b (b.x2 - b.x1) * (b.y2 - b.y1); return inter_area / (area_a area_b - inter_area 1e-6f); }分母加1e-6f是防止除零这是工程上的习惯做法虽然理论上面积不会为 0但浮点误差可能导致极小值。批量计算时我用一个std::vectorint记录哪些框已经被剔除避免重复计算std::vectorbool suppressed(n, false); std::vectorint keep; for (int i 0; i n; i) { if (suppressed[i]) continue; keep.push_back(i); for (int j i 1; j n; j) { if (suppressed[j]) continue; if (boxes[i].class_id ! boxes[j].class_id) continue; if (compute_iou(boxes[i], boxes[j]) iou_threshold) { suppressed[j] true; } } }这个双层循环看起来是 O(n²)但实际上因为大量框被提前剔除实际复杂度远低于理论值。25200 个框经过置信度过滤后通常只剩几百个再经过 NMS 后只剩几十个。3.4 pybind11 接口封装把 C 函数暴露给 Python代码很简洁#include pybind11/pybind11.h #include pybind11/numpy.h #include pybind11/stl.h namespace py pybind11; py::array_tfloat nms( py::array_tfloat boxes, py::array_tfloat scores, py::array_tint class_ids, float iou_threshold, float score_threshold ) { // 获取输入数组的 buffer auto boxes_buf boxes.request(); auto scores_buf scores.request(); auto class_buf class_ids.request(); int n scores_buf.shape[0]; float* boxes_ptr static_castfloat*(boxes_buf.ptr); float* scores_ptr static_castfloat*(scores_buf.ptr); int* class_ptr static_castint*(class_buf.ptr); // 构建 BBox 列表 std::vectorBBox bbox_list; bbox_list.reserve(n); for (int i 0; i n; i) { if (scores_ptr[i] score_threshold) continue; BBox b; b.x1 boxes_ptr[i * 4 0]; b.y1 boxes_ptr[i * 4 1]; b.x2 boxes_ptr[i * 4 2]; b.y2 boxes_ptr[i * 4 3]; b.score scores_ptr[i]; b.class_id class_ptr[i]; bbox_list.push_back(b); } // 执行 NMS std::vectorint keep nms_cpu(bbox_list, iou_threshold); // 构建输出 py::array_tfloat result({(int)keep.size(), 6}); auto result_buf result.request(); float* result_ptr static_castfloat*(result_buf.ptr); for (size_t i 0; i keep.size(); i) { const BBox b bbox_list[keep[i]]; result_ptr[i * 6 0] b.x1; result_ptr[i * 6 1] b.y1; result_ptr[i * 6 2] b.x2; result_ptr[i * 6 3] b.y2; result_ptr[i * 6 4] b.score; result_ptr[i * 6 5] static_castfloat(b.class_id); } return result; } PYBIND11_MODULE(fast_nms, m) { m.doc() Fast NMS implementation for RK3588; m.def(nms, nms, Non-Maximum Suppression, py::arg(boxes), py::arg(scores), py::arg(class_ids), py::arg(iou_threshold) 0.45f, py::arg(score_threshold) 0.25f); }这里有几个细节值得说第一py::array_tfloat默认不检查数组是否连续。如果 Python 传进来的是切片或者转置过的数组内存可能不连续直接按指针访问会出错。稳妥的做法是加py::array::c_style | py::array::forcecast标志让 pybind11 自动处理。第二request()返回的 buffer 在函数返回前一直有效不需要手动释放。第三输出数组我用的是[N, 6]格式前 4 列是坐标第 5 列是分数第 6 列是类别。这个格式和 YOLOv5 官方 Python 后处理的输出一致方便替换。3.5 编译与安装在 RK3588 上执行cd fast_nms_project python3 setup.py build_ext --inplace编译成功后会在当前目录生成fast_nms.cpython-38-aarch64-linux-gnu.so具体文件名取决于 Python 版本。把这个 .so 文件拷贝到你的项目目录就可以import fast_nms了。如果编译报错找不到 Python.h检查一下python3-dev是否安装sudo apt-get install python3-dev4. 性能对比与实测数据4.1 测试环境说明硬件RK3588 开发板8GB RAMCPU 频率 2.4GHz系统Ubuntu 20.04 aarch64Python3.8.10NumPy1.21.0测试数据YOLOv5s 在 640×640 输入下的原始输出25200 个候选框4.2 Python NMS 基准测试Python 版本我用的是 YOLOv5 官方 utils/general.py 里的non_max_suppression函数这是最常用的实现。测试代码import numpy as np import time from utils.general import non_max_suppression # 模拟 YOLOv5s 输出 prediction np.random.randn(1, 25200, 85).astype(np.float32) prediction torch.from_numpy(prediction) # 预热 for _ in range(5): _ non_max_suppression(prediction, conf_thres0.25, iou_thres0.45) # 计时 times [] for _ in range(50): start time.perf_counter() _ non_max_suppression(prediction, conf_thres0.25, iou_thres0.45) times.append(time.perf_counter() - start) print(fPython NMS 平均耗时: {np.mean(times)*1000:.2f} ms)实测结果平均 178.3ms最慢 195ms最快 172ms。这个数据和我预期一致Python 在这种不规则计算上确实吃力。4.3 C NMS 基准测试C 版本调用方式import numpy as np import time import fast_nms # 准备数据 boxes np.random.rand(25200, 4).astype(np.float32) * 640 scores np.random.rand(25200).astype(np.float32) class_ids np.random.randint(0, 80, 25200).astype(np.int32) # 预热 for _ in range(10): _ fast_nms.nms(boxes, scores, class_ids, 0.45, 0.25) # 计时 times [] for _ in range(100): start time.perf_counter() _ fast_nms.nms(boxes, scores, class_ids, 0.45, 0.25) times.append(time.perf_counter() - start) print(fC NMS 平均耗时: {np.mean(times)*1000:.3f} ms)实测结果平均 0.497ms最慢 0.62ms最快 0.43ms。4.4 加速比计算与结果分析加速比 178.3 / 0.497 ≈358.7 倍四舍五入就是 359 倍。这个数字看起来很夸张但拆解一下就很合理环节Python 耗时C 耗时加速比数据预处理15ms0.05ms300x排序45ms2.1ms21xIoU 计算100ms0.3ms333x结果构建18ms0.05ms360x总计178ms0.5ms356x排序环节加速比最低因为std::sort和 NumPy 的argsort都是 O(n log n)差距主要在常数项。IoU 计算加速比最高因为 C 的循环没有 Python 解释器开销而且编译器做了大量优化。实操心得如果你的场景对精度要求极高可以把-ffast-math去掉性能会下降约 15%但浮点结果更接近 IEEE 标准。我实测下来NMS 对这点精度差异完全不敏感所以保留了-ffast-math。5. 踩坑记录与常见问题排查5.1 编译期问题问题一pybind11/pybind11.h: No such file or directory原因pybind11 没安装或者 include 路径不对。解决pip3 install pybind11 python3 -m pybind11 --includes # 输出类似 -I/usr/include/python3.8 -I/usr/local/lib/python3.8/dist-packages/pybind11/include # 把这两个路径加到 setup.py 的 include_dirs 里问题二undefined symbol: _Py_Dealloc原因编译时链接的 Python 版本和运行时不一致。比如编译用的是 Python 3.8运行时用的是 Python 3.9。解决确认python3 --version和python3-config --includes输出的是同一个版本。如果有多个 Python用虚拟环境隔离。问题三-marchnative报错unknown architecture原因交叉编译时编译器不知道目标 CPU。解决改成-marcharmv8.2-a或者直接在 RK3588 上编译。5.2 运行期问题问题四TypeError: nms(): incompatible function arguments原因传入的 NumPy 数组 dtype 不对。pybind11 的py::array_tfloat要求 float32如果你传的是 float64 会报错。解决在 Python 侧显式转换boxes boxes.astype(np.float32) scores scores.astype(np.float32) class_ids class_ids.astype(np.int32)问题五结果和 Python 版本对不上原因可能是 IoU 阈值处理方式不同或者类别分组逻辑有差异。排查步骤先用同一组数据把 score_threshold 设为 0iou_threshold 设为 1.0看输出框数量是否一致逐步降低 iou_threshold对比每个阶段的输出检查 class_id 的处理Python 版本可能做了 class-wise NMS而 C 版本是全局 NMS我踩过的坑是Python 版本在计算 IoU 时用了torchvision.ops.nms它对坐标的格式要求是 (x1, y1, x2, y2)而我一开始传的是 (cx, cy, w, h)导致结果完全不对。后来统一转成角点坐标才对齐。问题六内存泄漏原因pybind11 的py::array_t如果手动管理指针容易忘记释放。解决尽量用request()返回的 buffer不要手动 new/delete。如果确实需要动态分配用std::vector或者智能指针。5.3 性能调优技巧技巧一预分配内存std::vector的reserve()能避免多次扩容。25200 个 BBox每个 24 字节总共约 600KB一次分配到位比反复扩容快 3-5 倍。技巧二用std::vectorbool做标记std::vectorbool是位压缩的比std::vectorchar省内存但访问速度稍慢。如果 N 很大100000用std::vectorchar更快如果 N 在几万级别std::vectorbool的缓存友好性更好。技巧三提前过滤低分框在构建 BBox 列表时就过滤掉 score 低于阈值的框能大幅减少后续计算量。实测 25200 个框经过 0.25 阈值过滤后通常只剩 300-500 个NMS 的实际计算量下降了 50 倍以上。技巧四类别分组并行如果 CPU 核心多可以把不同类别的 NMS 放到不同线程里并行执行。RK3588 有 8 个核理论上能再快 4-6 倍。不过线程创建和同步有开销如果框数量不多反而可能变慢。我实测下来单线程已经够快0.5ms没必要上多线程。5.4 常见问题速查表问题现象可能原因解决方法编译报错找不到头文件pybind11 未安装pip3 install pybind11运行时符号未定义Python 版本不一致统一编译和运行环境参数类型错误NumPy dtype 不对显式 astype 转换结果不一致坐标格式或 NMS 逻辑差异逐步对比中间结果性能不如预期编译优化未开启加 -O3 -marchnative内存占用高未预分配或未及时释放用 reserve 和 RAII6. 完整部署流程与集成建议6.1 从零到跑通的完整步骤假设你已经在 RK3588 上跑通了 RKNN 推理拿到了原始输出。完整集成步骤如下第一步创建项目目录mkdir -p ~/yolov5_rk3588/fast_nms cd ~/yolov5_rk3588/fast_nms第二步编写fast_nms.cpp内容就是前面第 3 节的核心代码。第三步编写setup.py内容见 3.1 节。第四步编译python3 setup.py build_ext --inplace第五步测试python3 -c import fast_nms; print(fast_nms.nms.__doc__)如果输出Non-Maximum Suppression说明编译成功。第六步集成到推理脚本import fast_nms import numpy as np def postprocess(outputs, conf_thres0.25, iou_thres0.45): # outputs 是 RKNN 推理的原始输出形状 [1, 25200, 85] # 拆分成 boxes, scores, class_ids boxes outputs[..., :4].reshape(-1, 4) scores outputs[..., 4].reshape(-1) class_scores outputs[..., 5:].reshape(-1, 80) class_ids np.argmax(class_scores, axis1) max_class_scores np.max(class_scores, axis1) final_scores scores * max_class_scores # 调用 C NMS result fast_nms.nms( boxes.astype(np.float32), final_scores.astype(np.float32), class_ids.astype(np.int32), iou_thres, conf_thres ) return result第七步跑一帧实际图像对比 Python 版本和 C 版本的输出框数量和坐标确认精度一致。6.2 与 RKNN 推理的对接细节RKNN 的输出格式和 PyTorch 略有不同。RK3588 的 NPU 输出通常是 NHWC 或者 NCHW需要根据模型转换时的配置来确定。我这边转换时用的是 NCHW输出形状是[1, 25200, 85]其中 85 4坐标 1置信度 80类别分数。坐标的格式也要注意YOLOv5 训练时用的是归一化的 (cx, cy, w, h)但 RKNN 转换后可能已经解码成像素坐标。我建议在 Python 侧统一转成 (x1, y1, x2, y2) 的像素坐标再传给 C NMS这样 C 代码不需要关心坐标格式。转换代码def xywh2xyxy(x): y np.copy(x) y[..., 0] x[..., 0] - x[..., 2] / 2 y[..., 1] x[..., 1] - x[..., 3] / 2 y[..., 2] x[..., 0] x[..., 2] / 2 y[..., 3] x[..., 1] x[..., 3] / 2 return y6.3 多线程推理场景下的注意事项如果你的 pipeline 是多线程的比如一个线程做推理一个线程做后处理需要注意 pybind11 的 GIL全局解释器锁。默认情况下C 函数执行时会持有 GIL阻塞其他 Python 线程。解决方法是在 C 函数里释放 GILpy::array_tfloat nms(...) { py::gil_scoped_release release; // 执行 NMS 计算 // ... py::gil_scoped_acquire acquire; // 构建返回值 // ... }这样在 NMS 计算期间其他 Python 线程可以继续执行。不过要注意释放 GIL 后不能访问任何 Python 对象所以数据拷贝要在释放前完成。我实测下来单线程场景下释放 GIL 没有性能提升反而有轻微开销。只有多线程场景才值得这么做。6.4 后续优化方向当前版本已经能满足大部分实时检测需求0.5ms 的 NMS 相对于 15ms 的推理可以忽略不计。如果还想进一步优化有几个方向第一用 NEON 指令手动向量化 IoU 计算。ARMv8 的 NEON 有 128 位寄存器一次能处理 4 个 float理论上能再快 2-3 倍。不过手写 NEON intrinsics 代码可读性差维护成本高除非性能极度敏感否则不建议。第二用 OpenMP 并行化外层循环。RK3588 有 8 个核如果框数量超过 1000并行化能带来明显收益。但要注意线程创建开销小规模数据反而变慢。第三把 NMS 直接集成到 RKNN 的 C API 里省去 Python 调用开销。不过这样会失去 Python 的灵活性调试也麻烦我倾向于保持现状。7. 一些个人体会这个项目做下来最大的感受是边缘计算场景下后处理优化的重要性不亚于模型本身。很多人花大量时间调模型结构、量化精度却忽略了后处理这个瓶颈。实际上一个 0.5ms 的 NMS 和一个 180ms 的 NMS直接决定了你的产品能不能做到实时。pybind11 这个工具我是越用越喜欢。它让 C 和 Python 的边界变得非常自然你可以把性能敏感的部分用 C 写把逻辑控制和业务代码用 Python 写两边各取所长。而且 pybind11 的编译配置很简单不像 Cython 那样需要学一套新语法。最后分享一个小技巧如果你在 RK3588 上编译时遇到奇怪的链接错误先检查python3-config --ldflags的输出确保链接的是正确的 Python 库。我踩过一次坑系统里有 Python 3.8 和 3.9 两个版本编译时用了 3.8 的头文件运行时却是 3.9报了一堆符号未定义的错误排查了半天才发现是版本不一致。用虚拟环境能彻底避免这个问题。