RK3588上YOLOv5后处理加速:pybind11实现C++ NMS,提速359倍

发布时间:2026/10/31 14:20:59
RK3588上YOLOv5后处理加速:pybind11实现C++ NMS,提速359倍 1. 从 Python 到 C为什么 NMS 成了整条链路的瓶颈在 RK3588 上把 YOLOv5s 跑起来推理部分其实没有想象中那么折腾。RKNN 工具链把模型转成 rknn 格式之后NPU 负责前向计算单帧推理时间稳定在 20ms 上下这个成绩对于一块嵌入式板子来说已经相当能打了。但真正让我头疼的是推理之后的那段后处理——尤其是 NMSNon-Maximum Suppression非极大值抑制。一开始我用的是官方 Python demo 里的后处理逻辑跑单张图片的时候感觉还行几百毫秒的事。可一旦切到视频流或者连续推理模式问题就暴露出来了每帧后处理耗时居然能到 300ms 以上比推理本身还慢十几倍。这就好比你去餐厅吃饭厨房出餐只要两分钟结果服务员摆盘摆了一刻钟整桌人等得干着急。我拿time.time()在 Python 里卡了一下时间发现大头全在 NMS 的循环上。YOLOv5s 在 640x640 输入下三个检测头加起来会输出 25200 个候选框80x80x3 40x40x3 20x20x3 25200。Python 里用嵌套循环去算 IoU、做排序、逐轮筛选每个框都要和剩下的框比一遍这个计算量在解释型语言里就是灾难。更别提 Python 的 GIL 锁让多线程也救不了场。所以这一篇的核心目标很明确把 NMS 从 Python 搬到 C用 pybind11 做成 Python 可调用的扩展模块让后处理耗时从 300ms 级别压到 1ms 以内。实测下来加速比稳定在 359 倍左右整条链路的帧率直接从 3 FPS 拉到了 25 FPS 以上。这个提升不是靠什么黑科技就是把该用编译型语言干的活交还给编译型语言。如果你也在 RK3588 或者类似边缘设备上部署 YOLO 系列模型并且被后处理拖了后腿那这套思路可以直接抄作业。不需要你精通 C只要会写基本的 Python能看懂循环和数组操作跟着走一遍就能落地。2. 方案选型为什么是 pybind11 而不是其他路子2.1 几种常见的加速方案对比把 Python 后处理换成更快的实现摆在面前的路其实有好几条。我当初把能想到的方案都列了一遍逐个评估之后才定的 pybind11。方案思路优点缺点适用场景纯 Python 优化用 NumPy 向量化替代循环改动小无需编译NMS 逻辑复杂向量化后仍有大量中间数组开销实测只快 2-3 倍对性能要求不高的场景Cython把 Python 代码编译成 C 扩展语法接近 Python需要写 .pyx 文件类型声明繁琐调试体验差已有大量 Python 代码想渐进加速Numba JIT运行时编译加速加个装饰器就能用首次编译慢对复杂控制流支持有限RK3588 上 LLVM 后端兼容性一般数值计算密集且逻辑简单的场景纯 C 可执行文件整个后处理独立成二进制性能最好和 Python 主流程割裂数据传递麻烦全 C 部署的项目pybind11写 C 扩展模块供 Python 调用性能接近纯 C接口自然编译配置简单需要写 C 代码要配编译环境Python 主流程 局部加速的混合架构我最终选 pybind11核心理由有三条。第一它生成的模块在 Python 里调用起来和普通函数没区别import之后直接传 NumPy 数组进去返回也是 NumPy 数组数据流转零摩擦。第二编译配置极其简单一个setup.py加几行Extension定义就完事不需要手写复杂的构建脚本。第三它对 NumPy 的支持是原生级别的py::array_tfloat可以直接映射到 C 的连续内存省掉了手动解析 buffer 的麻烦。2.2 pybind11 在 RK3588 上的适配要点RK3588 是 ARM64 架构跑的是 Linux 系统编译 C 扩展和 x86 上略有不同。我踩过的坑主要集中在编译器和 Python 头文件路径上。首先确认板子上的 Python 版本和开发头文件。我用的是系统自带的 Python 3.8需要装python3-dev包否则编译时会报Python.h: No such file or directory。命令很简单sudo apt update sudo apt install python3-dev python3-pip pip3 install pybind11 numpy然后确认 g 版本。RK3588 的官方系统一般自带 g 9 以上支持 C14 没问题。如果你的系统比较老可能需要手动升级。检查命令g --version编译的时候有个关键点-O3优化必须开-marchnative在 ARM 上不一定生效但-mcpucortex-a76.cortex-a55可以针对 RK3588 的大小核架构做优化。不过实测下来-O3加上-ffast-math已经能拿到绝大部分性能收益-mcpu的边际提升大概在 5% 左右可加可不加。注意RK3588 上编译 pybind11 扩展时如果遇到undefined symbol相关的错误大概率是 Python 链接库路径不对。可以在setup.py里显式指定library_dirs和include_dirs指向/usr/include/python3.8和/usr/lib/aarch64-linux-gnu。3. NMS 原理拆解从数学公式到代码逻辑3.1 IoU 计算NMS 的地基NMS 的核心就一件事找出那些“重叠太多”的框把分数低的干掉。判断重叠程度的指标就是 IoUIntersection over Union交并比。两个框 A 和 BIoU 的定义是IoU (A ∩ B) / (A ∪ B)分子是交集面积分母是并集面积。并集面积等于 A 的面积加 B 的面积减去交集面积。所以实际计算只需要三步算交集、算各自面积、套公式。交集的计算有个容易写错的地方。两个矩形框的交集左上角坐标取两者左上角的最大值右下角坐标取两者右下角的最小值。如果算出来的右下角坐标小于左上角坐标说明两个框根本不相交交集面积为 0。用代码表达就是float iou(const float* box_a, const float* box_b) { float x1 std::max(box_a[0], box_b[0]); float y1 std::max(box_a[1], box_b[1]); float x2 std::min(box_a[2], box_b[2]); float y2 std::min(box_a[3], box_b[3]); float inter_w std::max(0.0f, x2 - x1); float inter_h std::max(0.0f, y2 - y1); float inter_area inter_w * inter_h; float area_a (box_a[2] - box_a[0]) * (box_a[3] - box_a[1]); float area_b (box_b[2] - box_b[0]) * (box_b[3] - box_b[1]); return inter_area / (area_a area_b - inter_area 1e-6f); }那个1e-6f是防止除零的虽然理论上两个有效框的并集面积不会为零但加上更保险。3.2 标准 NMS 的算法流程标准 NMS 是逐类别做的流程如下按置信度分数从高到低排序所有候选框。取出分数最高的框加入保留列表。计算剩余所有框和这个框的 IoU把 IoU 超过阈值的框删掉。重复步骤 2 和 3直到候选框列表为空。这个流程的时间复杂度是 O(n²)n 是候选框数量。25200 个框的话最坏情况下要算 25200² 次 IoU大概是 6.3 亿次。Python 里每次 IoU 计算涉及多次函数调用和属性访问单次耗时可能到微秒级别乘起来就是几百毫秒。C 里同样的计算单次 IoU 就是几个浮点运算和比较CPU 一个周期能处理好几个单次耗时在纳秒级别。这就是 359 倍差距的来源——不是算法变了是执行效率变了。3.3 为什么不用 Soft-NMS 或者别的变体YOLOv5 官方用的就是标准 NMS我也没打算换。Soft-NMS 虽然在某些场景下能提升召回率但它需要多轮迭代衰减分数计算量比标准 NMS 还大。在 RK3588 这种算力有限的设备上后处理的目标是“够快够用”不是“刷榜”。标准 NMS 配合合适的置信度阈值和 IoU 阈值实际检测效果已经能满足绝大多数工业场景。4. C 实现细节从零写出可用的 NMS 模块4.1 数据结构设计C 里操作数据结构体设计直接影响缓存命中率和代码可读性。我定义了一个简单的Box结构struct Box { float x1, y1, x2, y2; float score; int class_id; };用std::vectorBox存储所有候选框。这里有个细节YOLOv5 的输出格式是[cx, cy, w, h, obj_conf, class_conf...]需要先转换成[x1, y1, x2, y2, score, class_id]。转换在 C 里做比在 Python 里做快得多因为涉及大量逐元素运算。转换公式x1 cx - w / 2 y1 cy - h / 2 x2 cx w / 2 y2 cy h / 2 score obj_conf * max(class_conf) class_id argmax(class_conf)注意 YOLOv5 的类别分数是独立 sigmoid 输出的不是 softmax。所以最终置信度是目标置信度乘以类别置信度的最大值。4.2 核心 NMS 函数实现完整的 NMS 函数签名如下std::vectorBox nms(std::vectorBox boxes, float iou_threshold, float conf_threshold);第一步是按置信度过滤。低于conf_threshold的框直接扔掉这一步能砍掉 90% 以上的候选框。25200 个框经过 0.25 的置信度阈值过滤后通常只剩几百到几千个。这个预处理极其重要它把 O(n²) 里的 n 从 25200 降到了几千计算量直接少两个数量级。std::vectorBox filtered; for (const auto box : boxes) { if (box.score conf_threshold) { filtered.push_back(box); } }然后是排序。用std::sort配合 lambda 表达式按分数降序排列std::sort(filtered.begin(), filtered.end(), [](const Box a, const Box b) { return a.score b.score; });接下来是逐类别 NMS。这里有个优化点可以先把框按类别分组每个类别单独做 NMS避免不同类别的框互相抑制。YOLOv5 的 80 类检测如果混在一起做 NMS一个高分的“人”可能会把旁边低分的“自行车”误删掉。std::vectorBox result; std::vectorbool suppressed(filtered.size(), false); for (size_t i 0; i filtered.size(); i) { if (suppressed[i]) continue; result.push_back(filtered[i]); for (size_t j i 1; j filtered.size(); j) { if (suppressed[j]) continue; if (filtered[i].class_id ! filtered[j].class_id) continue; if (iou(filtered[i].x1, filtered[j].x1) iou_threshold) { suppressed[j] true; } } }这个实现里suppressed数组避免了频繁的erase操作erase在 vector 中间删除元素会导致内存搬移代价很高。用标记法最后一次性收集结果效率更高。4.3 pybind11 绑定层写法C 函数写好后需要暴露给 Python。pybind11 的绑定代码很简洁#include pybind11/pybind11.h #include pybind11/numpy.h #include pybind11/stl.h namespace py pybind11; py::array_tfloat nms_binding(py::array_tfloat boxes, float iou_threshold, float conf_threshold) { auto buf boxes.request(); float* ptr static_castfloat*(buf.ptr); int num_boxes buf.shape[0]; int box_dim buf.shape[1]; std::vectorBox box_vec; for (int i 0; i num_boxes; i) { Box b; b.x1 ptr[i * box_dim 0]; b.y1 ptr[i * box_dim 1]; b.x2 ptr[i * box_dim 2]; b.y2 ptr[i * box_dim 3]; b.score ptr[i * box_dim 4]; b.class_id static_castint(ptr[i * box_dim 5]); box_vec.push_back(b); } auto result nms(box_vec, iou_threshold, conf_threshold); py::array_tfloat output({static_castint(result.size()), 6}); auto out_buf output.request(); float* out_ptr static_castfloat*(out_buf.ptr); for (size_t i 0; i result.size(); i) { out_ptr[i * 6 0] result[i].x1; out_ptr[i * 6 1] result[i].y1; out_ptr[i * 6 2] result[i].x2; out_ptr[i * 6 3] result[i].y2; out_ptr[i * 6 4] result[i].score; out_ptr[i * 6 5] static_castfloat(result[i].class_id); } return output; } PYBIND11_MODULE(nms_cpp, m) { m.doc() Fast NMS implementation for YOLOv5 on RK3588; m.def(nms, nms_binding, Run NMS on detection boxes, py::arg(boxes), py::arg(iou_threshold) 0.45, py::arg(conf_threshold) 0.25); }这里有几个关键点。py::array_tfloat默认要求输入是 C 连续的如果 Python 传进来的数组不是连续的pybind11 会自动拷贝一份这会带来额外开销。所以在 Python 侧调用前最好用np.ascontiguousarray()确保内存连续。输出数组的形状是(N, 6)每行是[x1, y1, x2, y2, score, class_id]。这个格式和 Python 侧的后续处理逻辑对齐拿到之后可以直接画框或者做其他操作。4.4 编译配置setup.py 怎么写编译用setup.py最省事from setuptools import setup, Extension import pybind11 ext_modules [ Extension( nms_cpp, [nms.cpp], include_dirs[ pybind11.get_include(), /usr/include/python3.8, ], languagec, extra_compile_args[-O3, -ffast-math, -stdc14], ), ] setup( namenms_cpp, version1.0, ext_modulesext_modules, )编译命令python3 setup.py build_ext --inplace成功的话会在当前目录生成nms_cpp.cpython-38-aarch64-linux-gnu.soPython 里直接import nms_cpp就能用。实操心得如果编译时报pybind11找不到检查pip3 show pybind11的路径确保get_include()返回的目录存在。RK3588 上有时候 pip 安装的包路径和系统 Python 路径不一致手动指定include_dirs更稳妥。5. 性能实测359 倍加速是怎么来的5.1 测试环境与基准设定测试在 RK3588 开发板上进行系统是 Ubuntu 20.04CPU 是 4 核 A76 4 核 A55内存 8GB。YOLOv5s 模型输入 640x640RKNN 推理输出三个检测头的原始数据。Python 版本的 NMS 用的是官方 demo 里的实现核心逻辑是 NumPy 数组操作加 Python 循环。C 版本就是上面写的 pybind11 扩展。测试方法同一组输入数据25200 个候选框分别跑 100 次取平均耗时。输入数据固定避免随机性影响。5.2 耗时对比数据实现方式平均耗时相对加速比Python NumPy312 ms1xPython NumPy优化版98 ms3.2xC pybind110.87 ms359xPython 优化版是把部分循环用 NumPy 向量化替代比如 IoU 计算批量做。但 NMS 的逐轮筛选逻辑很难完全向量化所以只提升了 3 倍左右。C 版本 0.87ms 的成绩主要归功于三点编译型语言的执行效率、置信度预过滤把候选框从 25200 降到几百、以及suppressed标记法避免了内存搬移。5.3 整条链路的帧率变化后处理从 312ms 降到 0.87ms 之后整条链路的耗时分布发生了质变阶段优化前优化后图像预处理8 ms8 msNPU 推理22 ms22 ms后处理NMS312 ms0.87 ms结果绘制5 ms5 ms总计347 ms35.87 ms帧率2.9 FPS27.9 FPS帧率从 2.9 拉到 27.9接近 10 倍提升。虽然后处理只占总时间的 90% 降到 2.4%但正是这 2.4% 的优化让整条链路从“不可用”变成了“流畅”。注意这个测试是在 CPU 上跑的。RK3588 的 NPU 推理是异步的实际部署时可以用多线程让推理和后处理重叠进一步压榨性能。但那是另一个话题了这里先把单帧延迟打下来。6. 踩坑记录与排查技巧6.1 编译期常见问题问题一Python.h找不到这是最常见的问题。RK3588 的官方系统镜像有时候不带 Python 开发头文件。解决方法sudo apt install python3-dev如果 apt 源里没有可以手动指定 Python 安装路径。用python3 -c import sysconfig; print(sysconfig.get_path(include))找到头文件目录加到include_dirs里。问题二链接时undefined reference to PyXXX这是 Python 库没链接上。在setup.py的Extension里加libraries[python3.8], library_dirs[/usr/lib/aarch64-linux-gnu],具体版本号根据你的 Python 版本调整。问题三编译出来的 .so 文件导入时报ImportError大概率是模块名和文件名不一致。PYBIND11_MODULE(nms_cpp, m)里的nms_cpp必须和Extension(nms_cpp, ...)里的名字完全一致包括大小写。6.2 运行期常见问题问题一结果和 Python 版本对不上NMS 的结果对排序稳定性敏感。如果两个框分数相同Python 的sorted和 C 的std::sort可能给出不同顺序导致保留的框不一样。解决方法在排序时加一个次要排序键比如按x1坐标排保证顺序确定。问题二IoU 阈值设多少合适YOLOv5 官方默认是 0.45。实际用下来0.45 到 0.5 之间比较平衡。设太低会漏检重叠目标设太高会保留太多重复框。如果检测目标密集可以适当提高到 0.5 到 0.6。问题三置信度阈值怎么调默认 0.25 是 YOLOv5 的通用值。如果你的场景误检多提到 0.4 到 0.5如果漏检多降到 0.1 到 0.2。这个值需要根据实际数据调没有万能公式。6.3 性能相关的坑坑一输入数组不是 C 连续的Python 侧如果对数组做了转置、切片等操作内存可能不连续。pybind11 会静默拷贝带来额外开销。养成习惯boxes np.ascontiguousarray(boxes, dtypenp.float32)坑二频繁创建销毁 vector如果每帧都重新分配std::vectorBox内存分配的开销会累积。可以在 C 侧用静态变量或者对象池复用内存。不过实测下来对于 25200 个框的规模vector 的分配开销在微秒级别影响不大。真正的大头是 IoU 计算本身。坑三编译优化没开-O3和-O0的性能差距能有 5 到 10 倍。一定要确认extra_compile_args里带了-O3。可以用objdump -d nms_cpp.so | head看看汇编如果全是函数调用没有内联说明优化没生效。7. 从 NMS 延伸还能加速什么NMS 搞定之后我把整条链路又过了一遍发现还有几个地方可以继续压榨。图像预处理YOLOv5 的 letterbox 缩放和归一化Python 里用 OpenCV 做大概 8ms。如果换成 C 实现能压到 2ms 以内。不过 8ms 相对于 35ms 的总耗时占比不算大优先级不高。结果绘制画框和标签在 Python 里用 OpenCV 大概 5ms。如果检测框数量不多这个开销可以接受。如果框特别多比如上百个可以考虑用 C 画好再传回 Python。多线程流水线RK3588 有 8 个核完全可以把预处理、推理、后处理分到不同线程用队列串起来。这样单帧延迟不变但吞吐量能翻倍。这个方案适合视频流场景单张图片检测用不上。模型量化YOLOv5s 的 RKNN 模型如果做 INT8 量化推理时间能从 22ms 降到 12ms 左右。但量化会带来精度损失需要评估是否可接受。这些优化点里NMS 的性价比最高因为它的耗时占比最大改动最局部收益最明显。其他的可以按需推进。8. 一些实操建议和代码组织方式C 代码建议单独放一个目录和 Python 主程序分开。我的目录结构是这样的project/ ├── cpp/ │ ├── nms.cpp │ └── setup.py ├── python/ │ ├── main.py │ ├── detector.py │ └── utils.py └── models/ └── yolov5s.rknn编译的时候在cpp/目录下执行python3 setup.py build_ext --inplace生成的.so文件留在cpp/里。Python 侧通过sys.path.append(cpp)把目录加进来然后import nms_cpp。如果要在多台设备上部署可以把编译好的.so文件直接拷贝过去只要 Python 版本和架构一致就能用。但更稳妥的做法是在目标设备上现场编译避免 ABI 不兼容的问题。Python 侧调用 NMS 的代码大概长这样import numpy as np import nms_cpp def postprocess(outputs, conf_thres0.25, iou_thres0.45): boxes decode_outputs(outputs) # 转成 (N, 6) 的 float32 数组 boxes np.ascontiguousarray(boxes, dtypenp.float32) result nms_cpp.nms(boxes, iou_thres, conf_thres) return resultdecode_outputs负责把 RKNN 的原始输出转成[x1, y1, x2, y2, score, class_id]格式。这部分逻辑如果也在 C 里做还能再快一点但 Python 里用 NumPy 向量化已经够用了。实操心得调试 C 扩展的时候可以在 C 里加printf输出中间结果编译后直接看终端输出。比用 gdb 单步调试方便得多。确认逻辑正确后再把 printf 删掉重新编译一版干净的。最后说一个容易被忽略的点NMS 的输入框数量在不同场景下差异很大。如果图像里目标很少候选框可能只有几百个C 的优势没那么明显。但如果场景复杂、目标密集候选框上万那 C 的加速效果就是决定性的。所以这套方案在工业质检、密集行人检测这类场景下价值最大。

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询