1. 为什么要在 RK3588 上死磕 NMS 加速

RK3588 这颗芯片在边缘计算圈子里火得一塌糊涂，8 核 CPU（4×A76 + 4×A55）加上 6 TOPS 的 NPU，跑 YOLOv5s 这种轻量级检测模型本来应该是轻轻松松的事。但真正把模型部署上去跑起来之后，很多人会发现一个很尴尬的现象：NPU 推理只花了十几毫秒，后处理却要花上百毫秒，整个 pipeline 的帧率被死死卡在个位数。这个瓶颈十有八九就出在 NMS（Non-Maximum Suppression，非极大值抑制）上。

我这次做的项目就是在 RK3588 上从零部署 YOLOv5s，前面几篇已经搞定了模型转换、RKNN 推理、Python 后处理验证。到了第五篇，核心目标很明确：把 Python 版本的 NMS 换成 C++ 实现，通过 pybind11 暴露给 Python 调用，最终实测加速 359 倍。这个数字不是标题党，是真实跑出来的对比数据——Python NMS 单帧约 180ms，C++ 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 有好几种方案，我逐一评估过：

我最终选 pybind11，理由很实在：第一，它是纯头文件库，不需要额外安装运行时，拷贝几个头文件就能用；第二，API 设计非常直观，py::array_t<float>直接映射 NumPy 数组，省去了大量手动转换代码；第三，RK3588 上交叉编译工具链对 C++11 支持完善，没有兼容性问题。

2.2 NMS 算法的核心逻辑再梳理

很多人写 NMS 就是照抄一段代码，但不知道为什么这么写。我先把逻辑讲透。

NMS 的输入是一堆候选框，每个框有坐标 (x1, y1, x2, y2)、置信度 score、类别 class_id。输出是经过筛选的框。核心步骤：

1. 按 score 从高到低排序
2. 取出当前最高分的框，保留
3. 计算这个框与剩余所有框的 IoU
4. 剔除 IoU 超过阈值的框（认为是同一个目标）
5. 重复 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.py：

from setuptools import setup, Extension import pybind11 ext_modules = [ Extension( 'fast_nms', ['fast_nms.cpp'], include_dirs=[pybind11.get_include()], language='c++', extra_compile_args=['-O3', '-march=native', '-ffast-math', '-std=c++11'], ), ] setup( name='fast_nms', ext_modules=ext_modules, )

这里几个编译参数很关键：-O3开启最高级别优化，-march=native让编译器针对 RK3588 的 Cortex-A76 生成最优指令，-ffast-math放宽浮点精度换取速度（NMS 对精度不敏感，可以放心用）。

注意：-march=native在交叉编译时会报错，因为编译器不知道目标 CPU。交叉编译时改成-march=armv8.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在 C++11 之后是内省排序（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::vector<int>记录哪些框已经被剔除，避免重复计算：

std::vector<bool> suppressed(n, false); std::vector<int> 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_t<float> nms( py::array_t<float> boxes, py::array_t<float> scores, py::array_t<int> 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_cast<float*>(boxes_buf.ptr); float* scores_ptr = static_cast<float*>(scores_buf.ptr); int* class_ptr = static_cast<int*>(class_buf.ptr); // 构建 BBox 列表 std::vector<BBox> 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::vector<int> keep = nms_cpu(bbox_list, iou_threshold); // 构建输出 py::array_t<float> result({(int)keep.size(), 6}); auto result_buf = result.request(); float* result_ptr = static_cast<float*>(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_cast<float>(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_t<float>默认不检查数组是否连续。如果 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-dev

4. 性能对比与实测数据

4.1 测试环境说明

• 硬件：RK3588 开发板，8GB RAM，CPU 频率 2.4GHz
• 系统：Ubuntu 20.04 aarch64
• Python：3.8.10
• NumPy：1.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_thres=0.25, iou_thres=0.45) # 计时 times = [] for _ in range(50): start = time.perf_counter() _ = non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45) times.append(time.perf_counter() - start) print(f"Python 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(f"C++ NMS 平均耗时: {np.mean(times)*1000:.3f} ms")

实测结果：平均 0.497ms，最慢 0.62ms，最快 0.43ms。

4.4 加速比计算与结果分析

加速比 = 178.3 / 0.497 ≈358.7 倍，四舍五入就是 359 倍。

这个数字看起来很夸张，但拆解一下就很合理：

排序环节加速比最低，因为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，用虚拟环境隔离。

问题三：-march=native报错unknown architecture

原因：交叉编译时编译器不知道目标 CPU。

解决：改成-march=armv8.2-a，或者直接在 RK3588 上编译。

5.2 运行期问题

问题四：TypeError: nms(): incompatible function arguments

原因：传入的 NumPy 数组 dtype 不对。pybind11 的py::array_t<float>要求 float32，如果你传的是 float64 会报错。

解决：在 Python 侧显式转换：

boxes = boxes.astype(np.float32) scores = scores.astype(np.float32) class_ids = class_ids.astype(np.int32)

问题五：结果和 Python 版本对不上

原因：可能是 IoU 阈值处理方式不同，或者类别分组逻辑有差异。

排查步骤：

1. 先用同一组数据，把 score_threshold 设为 0，iou_threshold 设为 1.0，看输出框数量是否一致
2. 逐步降低 iou_threshold，对比每个阶段的输出
3. 检查 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::vector<bool>做标记

std::vector<bool>是位压缩的，比std::vector<char>省内存，但访问速度稍慢。如果 N 很大（>100000），用std::vector<char>更快；如果 N 在几万级别，std::vector<bool>的缓存友好性更好。

技巧三：提前过滤低分框

在构建 BBox 列表时就过滤掉 score 低于阈值的框，能大幅减少后续计算量。实测 25200 个框经过 0.25 阈值过滤后通常只剩 300-500 个，NMS 的实际计算量下降了 50 倍以上。

技巧四：类别分组并行

如果 CPU 核心多，可以把不同类别的 NMS 放到不同线程里并行执行。RK3588 有 8 个核，理论上能再快 4-6 倍。不过线程创建和同步有开销，如果框数量不多，反而可能变慢。我实测下来，单线程已经够快（0.5ms），没必要上多线程。

5.4 常见问题速查表

6. 完整部署流程与集成建议

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_thres=0.25, iou_thres=0.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, axis=1) max_class_scores = np.max(class_scores, axis=1) 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 y

6.3 多线程推理场景下的注意事项

如果你的 pipeline 是多线程的（比如一个线程做推理，一个线程做后处理），需要注意 pybind11 的 GIL（全局解释器锁）。默认情况下，C++ 函数执行时会持有 GIL，阻塞其他 Python 线程。

解决方法是在 C++ 函数里释放 GIL：

py::array_t<float> 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，报了一堆符号未定义的错误，排查了半天才发现是版本不一致。用虚拟环境能彻底避免这个问题。