ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

RKNN上部署YOLOv8seg:从PyTorch到边缘端实例分割实战

RKNN上部署YOLOv8seg:从PyTorch到边缘端实例分割实战 1. 为什么要在RKNN上跑YOLOv8seg1.1 边缘端实例分割的真实需求先说清楚一件事YOLOv8seg不是普通的检测模型它在检测框的基础上多了一个分割分支能输出每个目标的像素级掩码。这意味着它不仅能告诉你“图里有一只猫”还能告诉你“这只猫的每一个像素在哪里”。在边缘计算场景里这个能力非常值钱——比如智能安防里要精确框出闯入区域的人体轮廓比如工业质检里要分割出产品表面的瑕疵区域比如农业无人机要识别并分割出病害叶片的具体范围。但问题来了边缘设备的算力通常很有限。你不可能在RK3588这种嵌入式板子上直接跑PyTorch原版模型推理速度会慢到无法接受。这时候RKNN就派上用场了。RKNN是瑞芯微的一套神经网络推理框架专门针对自家NPU做了深度优化能把模型推理速度提升几倍甚至十几倍。我实测过同一个YOLOv8seg模型在RK3588上跑PyTorch原版大概3-5 FPS转成RKNN之后能跑到25-30 FPS这个差距是质变的。所以如果你手头有RK3588、RK3568、RK3566这类板子又想做实例分割那这套流程是绕不过去的。1.2 零基础也能跑通的底气在哪很多人一听到“模型转换”“后处理”这些词就头大觉得这是算法工程师才能干的事。其实不是。整个流程的核心逻辑就三步把PyTorch模型导出成ONNX把ONNX转成RKNN然后在板子上写推理代码加后处理。每一步都有现成的工具和脚本你不需要自己写算子也不需要懂NPU的底层架构。我写这篇东西的目的就是让你拿着一个训练好的YOLOv8seg模型照着走一遍就能在RKNN上跑起来。中间会遇到什么坑、参数怎么设、后处理怎么写我都会掰开揉碎讲清楚。你不需要有深度学习部署经验但最好懂一点Python能看懂基本的numpy操作。2. 环境搭建与工具链选型2.1 硬件和软件的最低配置先列一下我用的环境你可以对照自己的情况调整项目我的配置最低要求说明开发机Ubuntu 20.04 x86_64Ubuntu 18.04用于模型转换目标板RK3588RK3568/RK3566/RK3588必须带NPUPython3.83.6-3.9版本太新可能不兼容RKNN-Toolkit21.6.01.4.0转换工具rknn-toolkit-lite21.6.0对应版本板端推理库PyTorch1.13.11.8导出ONNX用Ultralytics8.0.x8.0YOLOv8官方库这里有个关键点RKNN-Toolkit2的版本和板端runtime的版本必须匹配。我踩过一次坑开发机上用1.6.0转的模型板子上装的是1.4.0的runtime结果加载模型直接报错。所以装之前先确认板子上的librknnrt.so版本strings /usr/lib/librknnrt.so | grep version就能看到。2.2 RKNN-Toolkit2安装的坑官方推荐用conda建虚拟环境我建议你也这么做因为RKNN-Toolkit2对依赖版本比较挑剔。安装命令大概是这样conda create -n rknn python3.8 conda activate rknn pip install rknn-toolkit21.6.0 -i https://mirrors.aliyun.com/pypi/simple/注意如果你用的是x86开发机装的是rknn-toolkit2如果你直接在板子上转换不推荐板子算力不够那要装rknn-toolkit-lite2。这两个包不一样别搞混了。还有一个常见问题安装完之后from rknn.api import RKNN报错大概率是numpy版本冲突。RKNN-Toolkit2 1.6.0要求numpy1.24你如果装了numpy 2.x就会出问题。直接pip install numpy1.23.5降下来就行。2.3 板端环境准备板子上需要装rknn-toolkit-lite2这个包是专门给板端推理用的比完整版轻量很多。安装方式取决于你的板子系统# 如果是Debian/Ubuntu系统 pip install rknn-toolkit-lite21.6.0 # 如果是Buildroot系统可能需要交叉编译 # 建议直接用官方提供的wheel包另外板子上还需要OpenCV。如果你要做图像预处理和后处理可视化没有cv2会很麻烦。Buildroot系统里可能没有预装需要自己交叉编译或者用静态库。提示板端推理时尽量用rknn-toolkit-lite2而不是完整版rknn-toolkit2前者内存占用小很多启动也快。3. 从PyTorch到ONNX模型导出的关键细节3.1 YOLOv8seg的模型结构特点YOLOv8seg和YOLOv8det最大的区别在于分割头。检测头输出的是边界框回归和类别概率分割头输出的是原型掩码prototype masks和掩码系数mask coefficients。最终的分割结果 原型掩码 × 掩码系数再经过sigmoid和裁剪得到每个实例的掩码。这个结构对导出ONNX有影响。因为分割头的输出维度比检测头多你在写后处理的时候要同时处理两类输出。具体来说YOLOv8seg的ONNX模型通常有两个输出output0: [1, 116, 8400] —— 检测输出116 4bbox 80类别 32掩码系数output1: [1, 32, 160, 160] —— 原型掩码输出注意这里的116如果你的类别数不是80那这个数字会变。比如你训练的是自定义数据集只有3个类别那116就变成433239。这个一定要搞清楚后处理的时候维度对不上就是这里的问题。3.2 导出ONNX的实操步骤Ultralytics的库已经封装好了导出接口一行命令就能搞定from ultralytics import YOLO model YOLO(yolov8n-seg.pt) model.export(formatonnx, opset12, simplifyTrue, imgsz640)但这里有几个参数需要你根据实际情况调整opset建议用12兼容性最好。用11也行但有些算子可能不支持。simplify一定要开能去掉很多冗余算子RKNN转换时更顺畅。imgsz根据你的训练尺寸来通常是640。如果你训练时用了其他尺寸这里要对应改。导出完之后用Netron打开ONNX文件确认一下输入输出维度。输入应该是[1, 3, 640, 640]输出是上面说的两个。如果输出维度不对后面全白搭。3.3 导出时容易踩的坑第一个坑动态维度。Ultralytics默认导出的是动态batch但RKNN对动态维度的支持有限。你最好在导出时固定batch1或者导出后用onnxsim固定维度。我一般直接在export里加batch1参数。第二个坑opset版本太高。有人用opset17导出结果RKNN-Toolkit2 1.6.0不支持某些算子。建议不要超过12。第三个坑模型里有自定义算子。如果你在训练时改了YOLOv8的结构加了自定义层那导出ONNX时可能会失败。这种情况要么把自定义层去掉要么自己写ONNX算子映射。注意导出ONNX之后一定要用onnx.checker.check_model()验证一下模型完整性别等到转RKNN的时候才发现问题。4. ONNX转RKNN参数配置与量化策略4.1 RKNN转换的基本流程转换的核心代码不长但参数很关键from rknn.api import RKNN rknn RKNN(verboseTrue) # 配置 rknn.config( mean_values[[0, 0, 0]], std_values[[255, 255, 255]], target_platformrk3588, quantized_dtypeasymmetric_quantized-8, optimization_level3 ) # 加载ONNX ret rknn.load_onnx(modelyolov8n-seg.onnx) if ret ! 0: print(Load ONNX failed!) exit(ret) # 构建 ret rknn.build(do_quantizationTrue, datasetdataset.txt) if ret ! 0: print(Build failed!) exit(ret) # 导出 ret rknn.export_rknn(yolov8n-seg.rknn)这段代码里mean_values和std_values是归一化参数。YOLOv8训练时用的是0-1归一化所以mean0std255。如果你训练时用了其他归一化方式这里要对应改。4.2 量化数据集准备do_quantizationTrue表示做int8量化这能大幅提升推理速度但需要提供校准数据集。dataset.txt里每行是一张图片的路径建议准备200-500张有代表性的图片覆盖你的实际应用场景。量化数据集的选取有个原则分布要尽量接近实际推理时的输入。比如你做安防监控那就用监控截图做工业质检就用产线照片。别随便拿ImageNet的图凑数量化误差会很大。我试过用100张图和500张图做量化精度差异大概在1-2%左右。如果对精度要求高建议用500张以上。但注意量化图片越多转换时间越长500张大概要跑10-15分钟。4.3 量化精度调优技巧如果你发现量化后精度掉得厉害可以试试这几个方法第一调整optimization_level。默认是3可以降到2试试有时候能改善精度。第二混合量化。RKNN支持对某些层不做量化保持float16。你可以在config里加quantized_algorithmnormal和quantized_methodchannel这两个参数对精度影响比较大。第三换量化算法。RKNN-Toolkit2支持mmse和kl_divergence两种量化算法默认是normal。如果精度不理想可以试试mmse但转换时间会更长。第四检查是否有异常值。有时候输入数据里有个别像素值特别大或特别小会拉大量化范围。可以在预处理时做clip把像素值限制在合理范围内。提示量化后的模型精度损失在1-3%以内是可以接受的如果超过5%就要认真排查了。5. 板端推理与后处理全解析5.1 板端推理代码框架板端推理的核心是加载rknn模型然后循环执行推理。基本框架如下from rknnlite.api import RKNNLite import cv2 import numpy as np rknn_lite RKNNLite() ret rknn_lite.load_rknn(yolov8n-seg.rknn) ret rknn_lite.init_runtime(core_maskRKNNLite.NPU_CORE_0) img cv2.imread(test.jpg) img cv2.resize(img, (640, 640)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) outputs rknn_lite.inference(inputs[img])注意core_mask参数RK3588有三个NPU核心可以指定用哪个。如果你要跑多个模型可以分配到不同核心上并行。5.2 后处理的核心逻辑后处理是整个流程里最复杂的部分。YOLOv8seg的输出需要经过解码、NMS、掩码生成三个步骤。第一步解码检测输出output0的维度是[1, 116, 8400]其中8400是候选框数量116是每个框的信息。前4个是bbox的xywh接着80个是类别分数最后32个是掩码系数。output0 outputs[0].reshape(1, 116, 8400) output0 output0[0].transpose(1, 0) # [8400, 116] boxes output0[:, :4] scores output0[:, 4:84] mask_coeffs output0[:, 84:]第二步NMS过滤先按类别分数阈值过滤再做NMSclass_ids np.argmax(scores, axis1) confidences np.max(scores, axis1) keep confidences 0.25 boxes boxes[keep] class_ids class_ids[keep] confidences confidences[keep] mask_coeffs mask_coeffs[keep] # NMS indices cv2.dnn.NMSBoxes( boxes.tolist(), confidences.tolist(), 0.25, 0.45 )第三步生成掩码这是YOLOv8seg后处理的关键。output1是原型掩码维度[1, 32, 160, 160]。每个实例的掩码 mask_coeffs output1得到[160, 160]的掩码图然后resize到原图尺寸。proto outputs[1].reshape(32, 160, 160) masks mask_coeffs proto.reshape(32, -1) # [N, 25600] masks masks.reshape(-1, 160, 160) masks 1 / (1 np.exp(-masks)) # sigmoid然后根据bbox裁剪掩码再resize到原图大小。5.3 后处理中的性能优化后处理在CPU上跑如果候选框多NMS和掩码生成会很耗时。我实测过8400个候选框做NMS大概要20-30ms掩码生成要10-15ms。加起来快40ms了比NPU推理本身还慢。优化方法有几个第一在NPU推理前就把置信度阈值调高减少候选框数量。比如从0.25提到0.5候选框能少一半。第二用numpy的向量化操作代替循环。比如sigmoid计算用1/(1np.exp(-x))比循环快几十倍。第三掩码生成时只对保留的框做计算别对所有8400个框都算。第四如果板子支持可以用OpenCV的UMat做GPU加速但RK3588的GPU性能一般提升有限。注意后处理的时间要算进整体延迟里。如果你要求30FPS那推理后处理必须控制在33ms以内。6. 常见问题与排查技巧实录6.1 模型转换失败排查表报错信息可能原因解决方法Unsupported op typeONNX里有RKNN不支持的算子用onnxsim简化或换opset版本Quantization failed校准数据集有问题检查图片路径、格式、数量Shape mismatch输入维度不对确认ONNX输入是[1,3,640,640]Out of memory模型太大或量化数据太多减少校准图片数量或换小模型Version mismatchtoolkit和runtime版本不一致统一版本号6.2 推理结果异常的排查思路如果你发现推理结果不对比如框的位置偏了、掩码是乱的按这个顺序排查先看预处理。确认输入图片的resize方式、颜色通道顺序、归一化参数和训练时一致。YOLOv8训练时用的是RGB如果你传了BGR进去结果肯定不对。再看输出解码。确认output0的维度是[1, 116, 8400]如果维度不对说明模型导出有问题。另外确认bbox的解码方式YOLOv8用的是xywh格式不是xyxy。最后看后处理。NMS的阈值、掩码的sigmoid、resize的插值方式这些都会影响最终结果。建议先用一张简单图片测试比如纯色背景上放一个明显物体这样容易定位问题。6.3 性能不达标的优化方向如果推理速度达不到预期可以从这几个方面优化模型层面换更小的模型比如yolov8n-seg换成yolov8n-seg参数量少很多。或者降低输入分辨率从640降到416速度能提升一倍。转换层面确认用了int8量化别用float16。int8比float16快2-3倍。另外optimization_level设成3。推理层面用多核NPU并行RK3588有三个核心可以同时跑三个推理任务。但要注意内存带宽瓶颈三个核心同时跑不一定能线性提升。后处理层面把NMS和掩码生成放到NPU上做但这需要改RKNN模型结构比较复杂。简单的方法是降低候选框数量。6.4 我踩过的几个坑第一个坑量化数据集里混入了灰度图。RKNN量化时要求所有图片都是三通道如果有一张灰度图转换会报错。解决办法是预处理时统一转成RGB。第二个坑ONNX模型里有Dropout层。虽然推理时Dropout不生效但RKNN转换时会报错。导出时要用model.eval()切换到推理模式。第三个坑板端OpenCV版本太老。Buildroot自带的OpenCV可能是3.x版本不支持某些函数。建议自己交叉编译OpenCV 4.x。第四个坑NPU核心绑定。如果你不指定core_maskRKNN会自动选一个核心但有时候会选到正在忙的核心导致推理变慢。建议显式指定。第五个坑内存泄漏。如果你在循环里反复创建RKNNLite对象而不释放内存会一直涨。正确的做法是创建一次循环里只调inference。提示每次修改模型或参数后先用一张图片验证结果正确性再跑批量测试。别一上来就跑整个数据集浪费时间。7. 从能跑到好用工程化建议7.1 封装成可复用的推理类别把推理代码写成脚本建议封装成一个类方便复用class YOLOv8SegRKNN: def __init__(self, model_path, conf_thres0.25, iou_thres0.45): self.rknn RKNNLite() self.rknn.load_rknn(model_path) self.rknn.init_runtime(core_maskRKNNLite.NPU_CORE_0) self.conf_thres conf_thres self.iou_thres iou_thres def preprocess(self, img): # resize, normalize, transpose pass def inference(self, img): # rknn inference pass def postprocess(self, outputs): # decode, nms, mask generation pass def detect(self, img): input_data self.preprocess(img) outputs self.inference(input_data) results self.postprocess(outputs) return results这样你换模型、调参数都方便也容易集成到其他系统里。7.2 多线程与流水线设计如果你要处理视频流单线程推理肯定不够。建议用生产者-消费者模式一个线程负责读帧和预处理一个线程负责NPU推理一个线程负责后处理和显示。这样能充分利用CPU和NPU的并行能力。但要注意RKNNLite不是线程安全的。如果你在多线程里调同一个RKNNLite对象会出问题。解决办法是每个线程创建独立的RKNNLite实例或者加锁串行化。7.3 精度与速度的平衡策略实际项目里你需要在精度和速度之间做取舍。我的经验是如果对精度要求高用yolov8m-seg或yolov8l-seg输入640int8量化后处理用完整流程。RK3588上大概能跑10-15 FPS。如果对速度要求高用yolov8n-seg输入416int8量化后处理简化比如只做NMS不做掩码。能跑到40-50 FPS。如果两者都要那就用yolov8s-seg输入512这是比较平衡的选择。另外掩码生成其实很耗时如果你只需要检测框不需要掩码可以在后处理里跳过掩码生成速度能提升30%左右。7.4 模型更新与版本管理最后说一个容易被忽视的问题模型版本管理。你训练了新模型转成RKNN之后怎么知道板子上跑的是哪个版本我的做法是在RKNN模型文件名里加版本号和日期比如yolov8n-seg-v1.2-20240115.rknn。然后在代码里加一个版本检查启动时打印模型版本。这样出问题的时候能快速定位。另外建议保留每个版本的ONNX和RKNN文件别覆盖。有时候新模型效果不好要回滚到旧版本没有备份就麻烦了。提示RKNN模型不能跨平台使用。你在RK3588上转的模型不能直接在RK3568上跑。每个平台都要单独转换。8. 一些实战中的零碎经验8.1 输入尺寸的选择YOLOv8默认是640x640但RKNN对非正方形输入的支持也可以。如果你做的是宽屏视频比如1920x1080可以改成640x384这样能减少计算量速度提升大概20%。但要注意改输入尺寸后后处理里的resize参数也要对应改。8.2 类别数的处理如果你训练的是自定义数据集类别数不是80那output0的维度会变。比如3个类别维度就是[1, 39, 8400]。后处理代码里的切片索引要对应改别硬编码80。8.3 掩码阈值的调整掩码生成后通常要做一个二值化阈值一般是0.5。但这个阈值可以调调高掩码更干净但可能漏掉边缘调低掩码更完整但可能有噪声。根据你的应用场景调没有标准答案。8.4 板端温度控制RK3588跑满NPU的时候发热很大如果散热不好会降频。建议加散热片或者小风扇。我实测过不加散热跑10分钟NPU频率从1.0GHz降到0.8GHz速度掉20%。8.5 调试工具的使用RKNN-Toolkit2自带一个精度分析工具可以逐层对比量化前后的输出差异。如果你发现量化后精度掉得厉害用这个工具能快速定位是哪一层的问题。命令是rknn.accuracy_analysis()需要提供量化前后的模型和测试图片。这个工具我强烈建议你在转换阶段就用上别等到板端跑起来才发现精度不对。在开发机上花10分钟分析比在板子上折腾半天强。8.6 关于后处理的再思考后处理这块我后来发现一个更高效的写法把NMS和掩码生成合并到一个numpy操作里避免多次数据拷贝。具体做法是先用布尔索引一次性取出所有保留的框和掩码系数然后批量做矩阵乘法。这样比循环处理每个框快很多。另外sigmoid计算可以用查表法加速但精度会损失一点。如果你对速度要求极高可以试试。8.7 模型加密与保护如果你要把模型部署到客户设备上可能需要考虑模型保护。RKNN支持模型加密在export_rknn的时候加encryptTrue参数。但加密后的模型只能在特定设备上跑换设备要重新加密。这个功能我用过几次稳定性还可以但会增加一点加载时间。8.8 最后分享一个调试技巧当你遇到推理结果完全不对的时候先别急着改代码。用一张纯色图片比如全黑或全白跑一遍看看输出是什么。如果输出全是0或者全是NaN那说明预处理或模型加载有问题。如果输出有值但不对那再查后处理。这个方法帮我省了很多时间因为纯色图片的预期输出是确定的没有检测框很容易判断问题出在哪一环。整个流程走下来从模型转换到板端跑通快的话半天慢的话两三天。主要时间花在环境配置和后处理调试上。但一旦跑通后面换模型、调参数就很快了。希望这些经验能帮你少走点弯路。
返回列表