1. 为什么环境搭建这一步值得单独拿出来讲
很多人拿到 RK3588 开发板之后,第一反应是直接找现成的镜像烧进去跑 demo,能出画面就觉得大功告成。但真到要把 YOLOv5s 这种自定义模型塞进去跑的时候,问题就全冒出来了:PyTorch 版本对不上、ONNX 导出算子不支持、板端推理库和模型输入输出节点对不齐。我前后在 RK3588 上折腾过好几轮 YOLO 系列模型的部署,踩过的坑基本都集中在“环境”和“模型获取”这两个看似最没技术含量的环节上。
这篇是系列的第二部分,专门讲环境搭建与模型获取。第一部分聊的是整体链路设计和硬件选型思路,这一部分落地到具体操作:PC 端怎么搭一个干净的 PyTorch 训练/导出环境,YOLOv5s 的权重从哪来、怎么验证,ONNX 怎么导出、导出时哪些参数必须盯死,以及 RK3588 板端的基础运行环境怎么准备。整套流程走完,你会得到一个可以在 PC 上正常推理、并且已经导出为 ONNX 的 YOLOv5s 模型,为后续的 RKNN 转换和板端部署打好底子。
适合谁看:手上已经有 RK3588 开发板、想跑自己模型的人;做过深度学习训练但没接触过端侧部署的人;以及被 ONNX 导出各种报错折磨过的朋友。我尽量把每一步的“为什么”讲清楚,而不是只丢一堆命令让你复制。
2. 整体环境架构与方案选型思路
2.1 三段式环境划分:PC 训练端、转换端、板端
在端侧部署这件事上,最容易犯的错误就是把所有环境混在一台机器上。我的做法是明确分成三段:
- PC 训练/导出端:负责 YOLOv5s 的权重加载、验证、ONNX 导出。这台机器通常是带独立显卡的台式机或笔记本,装完整的 PyTorch 生态。
- 转换端:负责把 ONNX 转成 RKNN。这一步可以在 PC 上做(用 RKNN-Toolkit2),也可以在板子上做(用 RKNN-Toolkit-Lite2 或板端工具链)。我倾向于在 PC 上转,因为转换过程吃内存、需要反复调试,PC 上效率高得多。
- 板端运行端:RK3588 上只装推理运行时(RKNN Runtime)和必要的系统依赖,保持干净。
为什么要这么分?因为 PyTorch 的版本依赖非常“娇气”,而 RKNN 工具链对 Python 版本、NumPy 版本又有自己的一套要求。混在一起装,十有八九会出现“装了 RKNN 之后 PyTorch 导不出来了”或者“升级了 NumPy 之后 RKNN 转换报错”这类互相打架的问题。用 conda 或者 venv 把环境隔离,是最省心的做法。
2.2 为什么选 YOLOv5s 而不是 v8 或别的版本
热词里能看到不少人在搜 rk3588部署yolov8,YOLOv8 确实更新、精度也更好,但我在 RK3588 上第一轮实践还是选 YOLOv5s,原因很实际:
第一,YOLOv5 的 ONNX 导出链路最成熟,社区里针对 RKNN 的适配案例最多,遇到问题好查。第二,YOLOv5s 模型小(约 7.2M 参数,权重文件 14MB 左右),RK3588 的 NPU 算力(6 TOPS)跑起来帧率很舒服。第三,YOLOv8 的导出结构和后处理跟 v5 有差异,第一次上手容易在“输出节点对不上”上卡住。等你把 v5 这条链路跑通了,再换 v8 只是改几个导出参数的事。
提示:如果你明确只做 v8,那环境搭建部分基本通用,差异主要在导出脚本和后处理。建议先用 v5 把链路打通,再迁移。
2.3 版本锁定:整个链路最容易被忽视的关键
端侧部署最怕的就是“版本漂移”。我见过太多案例,同样的代码,今天能跑明天就报错,原因就是某个依赖自动升级了。所以我的原则是:所有关键依赖全部锁版本,并且记录下来。
下面这张表是我实测跑通的一套版本组合,你可以直接参考:
| 环节 | 组件 | 版本 | 说明 |
|---|---|---|---|
| PC 端 | Python | 3.8 | RKNN-Toolkit2 对 3.8 支持最好 |
| PC 端 | PyTorch | 1.13.1 | 与 YOLOv5 兼容稳定 |
| PC 端 | torchvision | 0.14.1 | 与 torch 1.13.1 配套 |
| PC 端 | ONNX | 1.12.0 | 导出算子支持完善 |
| PC 端 | onnxruntime | 1.14.0 | 用于验证 ONNX 推理 |
| 转换端 | RKNN-Toolkit2 | 1.5.2 | 支持 RK3588 |
| 板端 | RKNN Runtime | 1.5.2 | 与 Toolkit 版本对应 |
| 板端 | 系统 | Ubuntu 20.04 | 官方适配成熟 |
版本对应关系是硬约束,尤其是 RKNN-Toolkit2 和板端 Runtime 必须版本一致,否则会出现模型加载失败。这一点后面排查章节还会细说。
3. PC 端 PyTorch 环境搭建实操
3.1 用 conda 建一个干净的隔离环境
我强烈建议用 conda 而不是系统 Python,原因前面说了,隔离。命令不复杂:
conda create -n yolov5_rknn python=3.8 -y conda activate yolov5_rknn建完之后先确认 Python 版本:
python --version # 应该输出 Python 3.8.x这里有个小细节:如果你用的是 WSL(热词里 pytorch环境搭建wsl 搜索量很高),conda 装完之后记得把 WSL 的内存限制调一下。默认 WSL2 会吃掉宿主机一半内存,导出 ONNX 和后续转换时容易爆内存。在用户目录下建.wslconfig:
[wsl2] memory=16GB processors=8 swap=8GB改完wsl --shutdown重启生效。这个配置我实测下来,导出和转换过程稳定很多。
3.2 安装 PyTorch 与依赖:别直接 pip install torch
新手最容易犯的错是直接pip install torch,结果装了个最新版,跟 YOLOv5 的代码不兼容。正确做法是去 PyTorch 官网找对应 CUDA 版本的安装命令。如果你机器没有 NVIDIA 显卡,就装 CPU 版,导出 ONNX 用 CPU 完全够。
CUDA 11.7 对应的命令(有显卡的情况):
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 \ --extra-index-url https://download.pytorch.org/whl/cu117纯 CPU 的情况:
pip install torch==1.13.1 torchvision==0.14.1 \ --extra-index-url https://download.pytorch.org/whl/cpu装完验证一下:
import torch print(torch.__version__) print(torch.cuda.is_available()) # 有显卡应为 True然后装 YOLOv5 需要的其他依赖。这里注意,不要直接pip install -r requirements.txt,因为那个文件里会指定 torch 版本,可能把你刚装好的覆盖掉。我的做法是手动装核心的几个:
pip install numpy==1.23.5 opencv-python==4.7.0.72 \ pyyaml==6.0 tqdm==4.65.0 matplotlib==3.7.1 \ pandas==1.5.3 seaborn==0.12.2 onnx==1.12.0 \ onnxruntime==1.14.0NumPy 锁 1.23.5 是有原因的:RKNN-Toolkit2 对 NumPy 版本敏感,1.24 以上会有兼容性问题。这个坑我在转换阶段踩过,提前锁好省事。
3.3 获取 YOLOv5 源码与权重
YOLOv5 的源码直接从官方仓库拉,但要注意版本。我建议用 v6.2 或 v7.0 这两个 tag,它们的导出逻辑稳定,社区适配案例多。
git clone https://github.com/ultralytics/yolov5.git cd yolov5 git checkout v7.0权重文件yolov5s.pt有两个来源:一是官方 release 页面直接下载,二是用脚本自动下载。我倾向于手动下载放到weights/目录,因为自动下载有时候网络不稳,下到一半断了还不好排查。
mkdir -p weights # 把下载好的 yolov5s.pt 放到 weights/ 目录 ls -lh weights/yolov5s.pt # 正常应该是 14MB 左右拿到权重后,第一件事是验证它能正常推理,别急着导出。用官方自带的检测脚本跑一张图:
python detect.py --weights weights/yolov5s.pt \ --source data/images/bus.jpg --device cpu跑完在runs/detect/exp/下能看到带框的结果图。这一步能过,说明 PyTorch 环境和权重都没问题。如果这一步就报错,先别往下走,把环境问题解决干净。
注意:detect.py 第一次运行会自动下载一些字体文件,如果网络不通会卡住。可以提前把 Arial.ttf 放到缓存目录,或者忽略这个警告,不影响推理。
3.4 环境自检清单
在进入模型导出之前,我习惯跑一个自检脚本,把关键依赖的版本都打出来,存成文件备查:
import torch, torchvision, numpy, onnx, onnxruntime, cv2 print("torch:", torch.__version__) print("torchvision:", torchvision.__version__) print("numpy:", numpy.__version__) print("onnx:", onnx.__version__) print("onnxruntime:", onnxruntime.__version__) print("opencv:", cv2.__version__) print("cuda available:", torch.cuda.is_available())把输出存成env_check.txt。后面如果出问题,对照这个文件能快速定位是不是版本变了。这个习惯帮我省了无数次排查时间。
4. YOLOv5s 模型获取与 ONNX 导出全流程
4.1 导出前的模型结构确认
导出 ONNX 之前,得先搞清楚 YOLOv5s 的输出结构。很多人导出后板端跑不对,就是因为不知道模型到底输出了什么。YOLOv5 默认导出三个输出头,对应三个不同尺度的特征图,形状分别是:
- 小目标:
[1, 255, 80, 80] - 中目标:
[1, 255, 40, 40] - 大目标:
[1, 255, 20, 20]
这里的 255 是3 * (5 + 80),即 3 个 anchor、每个 anchor 5 个框参数(xywh + obj)、80 个类别。输入是[1, 3, 640, 640]。
为什么要在导出前确认这个?因为 RKNN 转换时需要指定输入输出节点,如果你不知道输出是什么形状,后面根本没法配。而且 YOLOv5 的导出脚本有个--include onnx参数,不同版本导出的输出节点命名可能不一样,提前确认能避免踩坑。
4.2 导出 ONNX 的关键参数逐个拆解
YOLOv5 自带的export.py就能导出 ONNX,但参数得配对。我常用的命令:
python export.py --weights weights/yolov5s.pt \ --include onnx \ --imgsz 640 640 \ --batch-size 1 \ --opset 12 \ --device cpu \ --simplify逐个说这些参数为什么这么设:
--imgsz 640 640:输入分辨率。RK3588 的 NPU 对 640 这个尺寸支持很好,而且 640 是 32 的倍数,YOLOv5 的下采样要求输入能被 32 整除。如果你想用 320 或 416,也可以,但精度会降。
--batch-size 1:端侧推理基本都是单张,batch 设 1 最省内存,也避免动态 batch 带来的转换麻烦。
--opset 12:ONNX 算子集版本。这个很关键。opset 太高(比如 17)会引入 RKNN 不支持的算子,太低(比如 9)又可能缺一些必要算子。opset 12 是我实测在 RK3588 上兼容性最好的版本。
--device cpu:导出用 CPU 就行,用 GPU 导出有时候会有算子差异。
--simplify:调用 onnx-simplifier 做图优化,去掉冗余节点。这个强烈建议开,能显著减少转换时的算子问题。需要先装:
pip install onnx-simplifier==0.4.17导出成功后,会在weights/目录下生成yolov5s.onnx,大小约 28MB。
4.3 导出后的 ONNX 验证:别跳过这一步
导出完直接拿去转换,是另一个常见错误。ONNX 文件可能“导出成功但推理错误”,必须验证。我用 onnxruntime 跑一遍,跟 PyTorch 的结果对比:
import onnxruntime as ort import numpy as np import torch # 加载 ONNX sess = ort.InferenceSession("weights/yolov5s.onnx", providers=["CPUExecutionProvider"]) input_name = sess.get_inputs()[0].name print("input name:", input_name) print("input shape:", sess.get_inputs()[0].shape) for out in sess.get_outputs(): print("output:", out.name, out.shape) # 造一个随机输入测试 dummy = np.random.randn(1, 3, 640, 640).astype(np.float32) outputs = sess.run(None, {input_name: dummy}) for i, o in enumerate(outputs): print(f"output {i} shape:", o.shape)正常应该打印出三个输出,形状跟前面说的一致。如果输出数量不对,或者形状是动态的(显示为字符串而不是数字),说明导出参数有问题,得回去检查。
提示:如果输出形状里出现
-1或'batch'这种动态维度,转换 RKNN 时可能报错。解决办法是导出时确保 batch 固定为 1,并且不要开 dynamic axes。
4.4 用真实图片做端到端验证
随机输入只能验证形状,验证不了数值正确性。我还会用一张真实图片,分别跑 PyTorch 和 ONNX,对比输出差异:
import cv2 import numpy as np import torch import onnxruntime as ort img = cv2.imread("data/images/bus.jpg") img = cv2.resize(img, (640, 640)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.transpose(2, 0, 1)[None].astype(np.float32) / 255.0 # PyTorch 推理 model = torch.load("weights/yolov5s.pt", map_location="cpu")["model"] model.eval() with torch.no_grad(): pt_out = model(torch.from_numpy(img))[0].numpy() # ONNX 推理 sess = ort.InferenceSession("weights/yolov5s.onnx", providers=["CPUExecutionProvider"]) onnx_out = sess.run(None, {sess.get_inputs()[0].name: img})[0] # 对比最大差异 diff = np.abs(pt_out - onnx_out).max() print("max diff:", diff)正常情况下 max diff 应该在 1e-4 到 1e-3 量级。如果差异很大(比如 0.1 以上),说明导出过程有问题,可能是某个算子被错误转换了。这个验证步骤花不了几分钟,但能帮你提前发现大问题。
4.5 模型轻量化的一些前置考虑
热词里 yolov5s模型轻量化 也是高频搜索。如果你后续打算做量化(INT8),在导出阶段就要留个心眼。INT8 量化需要校准数据集,而校准数据的预处理方式必须和推理时完全一致。我的建议是:
- 导出 ONNX 时就把预处理(归一化、resize)固定下来,别在板端再做额外处理。
- 准备 100 到 200 张有代表性的校准图片,覆盖你的实际场景。
- 校准图片的尺寸、通道顺序、归一化方式,跟推理输入严格对齐。
这些在导出阶段想清楚,后面量化会顺很多。量化本身是下一部分的内容,这里先埋个伏笔。
5. RK3588 板端基础环境准备
5.1 系统选择与烧写要点
RK3588 板端系统我推荐 Ubuntu 20.04,官方适配最成熟,RKNN Runtime 的安装包也最全。烧写用瑞芯微的烧录工具,注意几个点:
第一,烧写前确认板子的启动模式拨码开关在正确位置(通常是 MaskROM 模式)。第二,USB 线要用质量好的,劣质线会导致烧写中途断开。第三,烧写完成后第一次启动会比较慢,耐心等,别以为卡死了就断电。
系统起来之后,先更新一下基础包:
sudo apt update sudo apt install -y python3-pip python3-dev cmake git \ libopencv-dev python3-opencv5.2 板端 Python 环境与 RKNN Runtime 安装
板端的 Python 环境我建议直接用系统自带的 Python 3.8,不要额外建 conda 环境,因为 RKNN Runtime 的 wheel 包对系统 Python 支持最好。
RKNN Runtime 的安装分两部分:一是rknn_toolkit_lite2(Python 接口),二是底层的librknnrt.so(C 库)。安装包从官方仓库获取,注意版本要和 PC 端 Toolkit2 一致。
# 安装 Python 接口 pip3 install rknn_toolkit_lite2-1.5.2-cp38-cp38-linux_aarch64.whl # 底层库通常随系统镜像已包含,确认一下 ls /usr/lib/librknnrt.so如果librknnrt.so不存在,需要手动拷贝到/usr/lib/并执行sudo ldconfig。
验证安装:
from rknnlite.api import RKNNLite rknn = RKNNLite() print("RKNN Lite init ok")能打印出来就说明板端环境 OK 了。
5.3 NPU 驱动版本确认
RK3588 的 NPU 驱动版本会影响模型加载。查看驱动版本:
cat /sys/kernel/debug/rknpu/version输出的版本号要跟 RKNN Runtime 匹配。如果驱动太旧,加载新版本转换的模型会失败。升级驱动通常需要重新烧写系统或者单独更新内核模块,比较麻烦,所以一开始就选对系统版本很重要。
注意:NPU 驱动版本和 RKNN Runtime 版本不匹配是板端部署最常见的失败原因之一。转换模型前,先确认这两个版本,能省掉大量排查时间。
5.4 板端推理库的依赖检查
RKNN Runtime 依赖一些系统库,缺了会报libxxx.so not found。我习惯用ldd检查:
ldd /usr/lib/librknnrt.so | grep "not found"如果输出为空,说明依赖齐全。如果有缺失,根据提示装对应的库。常见的缺失是libgomp和libstdc++,一般apt install就能解决。
6. 常见问题与排查技巧实录
6.1 ONNX 导出报错速查表
导出阶段的问题最集中,我整理了一张速查表:
| 报错信息 | 原因 | 解决办法 |
|---|---|---|
Unsupported ONNX opset | opset 版本过高 | 降到 12 |
Export failed: dynamic shape | 开了 dynamic axes | 关闭,固定 batch=1 |
ModuleNotFoundError: onnx | 没装 onnx | pip install onnx==1.12.0 |
simplify failed | onnx-simplifier 版本不匹配 | 装 0.4.17 |
| 输出节点数量不对 | 导出脚本版本问题 | 换 v7.0 tag |
RuntimeError: CUDA out of memory | 用 GPU 导出 | 加 --device cpu |
6.2 版本冲突的排查思路
版本冲突是最难查的问题,因为报错信息往往指向别处。我的排查顺序是:
第一步,跑前面那个环境自检脚本,把所有版本打出来。第二步,对照官方文档的版本要求,看哪个对不上。第三步,用pip list看有没有重复安装的包(比如同时装了 CPU 版和 GPU 版 torch)。第四步,如果还找不到,新建一个干净环境重装,往往比在原环境里修更快。
我踩过最坑的一次是 NumPy 版本问题:RKNN-Toolkit2 在 NumPy 1.24 上会报一个很隐晦的AttributeError,查了半天才发现是版本问题。所以前面强调锁 NumPy 1.23.5。
6.3 板端模型加载失败的排查
板端加载 RKNN 模型失败,常见原因和排查方法:
- 版本不匹配:确认 Toolkit2 和 Runtime 版本一致。
- 模型文件损坏:重新转换一次,或者用 md5 校验文件完整性。
- 输入输出节点配置错误:检查转换时的 config,输入名字和形状要对。
- 内存不足:RK3588 内存有限,大模型可能加载失败,用
free -h看内存。
排查时打开 RKNN 的日志:
rknn = RKNNLite(verbose=True)verbose 打开后会打印详细的加载过程,哪一步失败一目了然。
6.4 几个我踩过的坑
第一个坑:导出 ONNX 时忘了--simplify,结果转换 RKNN 时报了一堆算子不支持。加上 simplify 之后问题消失。所以这个参数别省。
第二个坑:用 GPU 导出 ONNX,结果某些算子在 GPU 和 CPU 上的实现有细微差异,导致板端结果和 PC 端对不上。后来统一用 CPU 导出,问题解决。
第三个坑:板端系统自带的 OpenCV 版本太老,读图后颜色通道顺序跟预期不一致,导致检测框位置偏移。解决办法是统一用cv2.cvtColor显式转换,别依赖默认行为。
第四个坑:校准数据集和推理输入预处理不一致,量化后精度掉得厉害。这个在量化阶段才暴露,但根子在导出阶段没统一预处理。所以前面强调预处理要固定。
7. 环境搭建完成后的下一步衔接
走到这里,你应该已经拿到了一个验证过的yolov5s.onnx,PC 端能正常推理,板端 RKNN Runtime 也装好了。接下来就是把这个 ONNX 转成 RKNN 模型,然后在板端跑起来。转换阶段的核心是 config 配置和量化校准,这部分坑也不少,尤其是 INT8 量化的精度问题。
我个人在实际操作中的体会是,环境搭建和模型获取这两步看似简单,但决定了后面 80% 的顺利程度。版本锁死、导出验证、板端驱动确认,这三件事做到位,后面的转换和部署会顺畅很多。如果这三步偷懒,后面会花几倍的时间去排查各种莫名其妙的问题。
最后分享一个小技巧:把整个环境搭建过程写成一个 shell 脚本,包括 conda 环境创建、依赖安装、版本检查。下次换机器或者重装系统,一条命令就能恢复环境,比手动一步步装靠谱得多。这个脚本我用了大半年,帮我省了不知道多少重复劳动。