RK3588平台YOLOv5模型训练与部署实战

在智能摄像头、工业质检和边缘安防等场景中，如何将高效的AI目标检测能力落地到国产芯片上，是当前许多开发者关注的焦点。RK3588作为瑞芯微推出的旗舰级SoC，集成了高达6TOPS算力的NPU，为本地化视觉推理提供了强大支撑。然而，从PyTorch模型训练到最终在设备端稳定运行，并非一键导出那么简单——尤其是面对YOLOv5这类主流但结构复杂的模型时，激活函数兼容性、格式转换异常、推理性能瓶颈等问题常常让人止步不前。

本文将以实际工程视角出发，带你完整走通一条基于RK3588的YOLOv5端到端开发路径：从环境搭建开始，经过数据准备、模型训练、ONNX优化，再到RKNN转换与板端部署，每一步都结合真实踩坑经验给出解决方案。我们特别推荐使用 Miniconda-Python3.9镜像 构建轻量可复现的开发环境，避免依赖冲突导致的“本地能跑，换机就崩”问题。

---

整个流程中最关键的一环，其实是开发环境的设计逻辑。很多团队在项目初期直接用系统Python安装一堆包，结果不同项目间版本打架，后期难以迁移。而Miniconda的优势在于它像一个“沙盒”，每个项目独立运行，互不影响。相比Anaconda动辄几百MB的体积，Miniconda仅包含conda和Python解释器，启动快、占用小，非常适合AI开发这种对依赖管理要求高的场景。

以当前主流框架为例，PyTorch 1.12+对ONNX的支持更完善，但又要求Python ≥3.7且 ≤3.10；CUDA 11.7对应的torch版本也必须精确匹配。如果不用虚拟环境，很容易出现pip install torch后发现装错了版本的情况。因此，第一步永远是创建干净的虚拟环境：

conda create -n yolov5-rk3588 python=3.9 -y
conda activate yolov5-rk3588

激活成功后终端会显示 (yolov5-rk3588) 前缀，说明你现在处于隔离环境中。接下来可以安全地安装核心依赖：

# 安装支持CUDA 11.7的PyTorch
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117

# 克隆适配RK3588的YOLOv5分支（重要！）
git clone https://github.com/airockchip/yolov5.git
cd yolov5
pip install -r requirements.txt

这里有个极易被忽视的关键点：不能使用Ultralytics官方的YOLOv5仓库。原因在于其主干网络广泛使用的SiLU（Swish）激活函数，在RK3588的NPU上并不受原生支持，会导致后续转换时报错 Unsupported layer: Swish。Rockchip官方维护的 airockchip/yolov5 分支已将这些层替换为ReLU，同时内置了针对RKNN Toolkit的预处理脚本，极大提升了部署成功率。

如果你有GPU可用，建议通过以下命令验证CUDA是否正常工作：

import torch
print(torch.cuda.is_available())  # 应返回 True
print(torch.__version__)

若返回False，请检查驱动或重装PyTorch。训练阶段启用GPU可显著缩短迭代周期，尤其是在处理大规模数据集时。

---

进入模型训练前，首先要解决的是数据问题。没有高质量标注，再强的架构也只是空中楼阁。我们推荐使用图形化工具 LabelImg 进行手动标注：

pip install labelimg
labelimg

打开工具后，设置图像目录（如 datasets/images/），并将标签保存路径指向 datasets/labels/。注意选择 PascalVOC (.xml) 格式输出，这是目前最通用的标注标准之一。

每张图片对应一个.xml文件，记录目标类别和边界框坐标。例如：

<annotation>
    <filename>car_001.jpg</filename>
    <size>
        <width>640</width>
        <height>480</height>
    </size>
    <object>
        <name>car</name>
        <bndbox>
            <xmin>100</xmin>
            <ymin>80</ymin>
            <xmax>300</xmax>
            <ymax>200</ymax>
        </bndbox>
    </object>
</annotation>

但YOLOv5训练需要的是归一化的 .txt 标签格式：class_id center_x center_y width height。为此，我们可以编写一个自动转换脚本：

import os
import xml.etree.ElementTree as ET

classes = ['car', 'person', 'bike']  # 自定义类别列表

def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    return x*dw, y*dh, w*dw, h*dh

def convert_annotation(image_id):
    in_file = open(f'datasets/annotations/{image_id}.xml', encoding='utf-8')
    out_file = open(f'datasets/labels/{image_id}.txt', 'w')

    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls not in classes:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
             float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

# 创建必要目录
os.makedirs('datasets/images/train', exist_ok=True)
os.makedirs('datasets/images/val', exist_ok=True)
os.makedirs('datasets/labels/train', exist_ok=True)
os.makedirs('datasets/labels/val', exist_ok=True)

# 批量转换所有XML文件
for file in os.listdir('datasets/annotations'):
    if file.endswith('.xml'):
        name = os.path.splitext(file)[0]
        convert_annotation(name)

执行后，所有标注将转为YOLO所需格式。记得按比例划分训练集与验证集（建议8:2），否则容易过拟合。

---

接下来是模型配置环节。按照YOLOv5规范组织目录结构：

yolov5/
├── data/
│   └── custom.yaml
├── datasets/
│   ├── images/
│   │   ├── train/
│   │   └── val/
│   └── labels/
│       ├── train/
│       └── val/

其中 custom.yaml 内容如下：

train: ../datasets/images/train
val: ../datasets/images/val
nc: 3
names: ['car', 'person', 'bike']

一切就绪后即可启动训练：

python train.py \
--img 640 \
--batch 16 \
--epochs 100 \
--data data/custom.yaml \
--weights yolov5s.pt \
--cfg models/yolov5s.yaml \
--name yolov5s_rk3588

几个实用提示：
- 若显存不足，可降低--batch至8或改用yolov5n；
- 训练过程中的loss曲线和mAP可在runs/train/yolov5s_rk3588/中查看；
- 最终生成的最优权重为best.pt，可用于后续导出。

---

当得到满意的best.pt模型后，下一步就是将其部署到RK3588上。由于该平台使用专有的 RKNN Toolkit 推理引擎，我们必须先将模型转为.rknn格式。

首先导出ONNX中间表示：

python export.py \
--weights runs/train/yolov5s_rk3588/best.pt \
--include onnx \
--img 640 \
--batch 1

务必注意：
- 输入尺寸固定为静态shape（batch=1）；
- ONNX Opset应不低于11；
- 若遇到不支持的操作符（如Slice），需在代码中手动替换为Split等等价结构。

然后在RK3588开发板或Linux主机上安装RKNN Toolkit2：

pip install rknn_toolkit_lite2

编写转换脚本 convert_rknn.py：

from rknn.api import RKNN

rknn = RKNN()

# 设置输入归一化参数
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3588')

ret = rknn.load_onnx(model='best.onnx')
if ret != 0:
    print('Load ONNX failed!')
    exit(ret)

# 构建模型（若需INT8量化，设do_quantization=True并提供校准集）
ret = rknn.build(do_quantization=False)
if ret != 0:
    print('Build RKNN failed!')
    exit(ret)

# 导出最终模型
rknn.export_rknn('yolov5s.rknn')
rknn.release()

运行脚本即可生成yolov5s.rknn，准备好推送到设备端。

---

部署阶段通常通过SSH连接RK3588开发板完成：

ssh firefly@192.168.1.100
# 默认密码通常是 fa

上传模型文件：

scp yolov5s.rknn firefly@192.168.1.100:/home/firefly/models/

在开发板上编写推理脚本 infer.py：

import cv2
import numpy as np
from rknnlite.api import RKNNLite

rknn = RKNNLite()
ret = rknn.load_rknn('yolov5s.rknn')
if ret != 0:
    print('Load RKNN failed!')
    exit(ret)

ret = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
if ret != 0:
    print('Init runtime failed!')
    exit(ret)

# 图像预处理
img = cv2.imread('test.jpg')
img = cv2.resize(img, (640, 640))
input_data = img[np.newaxis, :, :, ::-1].astype(np.float32) / 255.0

# 执行推理
outputs = rknn.inference(inputs=[input_data])
print("Output shape:", [o.shape for o in outputs])

# 简单后处理示意
pred = outputs[0]
# 此处可接入解码逻辑：过滤低置信度框、NMS去重等

print("Inference completed.")
rknn.release()

运行后预期输出类似：

Output shape: [(1, 25200, 85)]
Inference completed.

这表明模型已在NPU上成功加载并完成一次前向传播。

---

在实际调试过程中，以下几个问题是高频出现的：

问题	原因	解决方案
Unsupported layer: Swish	使用了SiLU激活函数	改用Rockchip分支或手动替换为ReLU
ONNX导出失败	动态shape或不支持op	固定输入尺寸，禁用自动形状推断
推理速度慢	未启用多核NPU	设置core_mask=RKNNLite.NPU_CORE_0_1_2
mAP偏低	数据量少或过拟合	加强数据增强、调整学习率策略

一些进阶优化建议值得尝试：
- INT8量化：在开启do_quantization=True的前提下，使用少量校准图像（约100~200张）进行校准，推理速度可提升2~3倍；
- 分辨率权衡：将输入从640×640降至416×416，虽牺牲部分精度，但帧率明显上升，适合实时性要求高的场景；
- 流水线设计：结合OpenCV多线程读图与NPU异步推理，实现“边取图、边推理”，提高整体吞吐量。

---

这条从训练到部署的完整链路，不仅适用于YOLOv5，也为其他模型迁移到RK3588平台提供了参考模板。更重要的是，它强调了一种工程思维：环境可复现、流程可追踪、问题可定位。当你不再被“为什么昨天还能跑今天就不行”困扰时，才能真正专注于算法本身的价值创造。

未来你可以在此基础上拓展更多应用，比如接入RTSP视频流做连续检测，或是结合GPIO控制实现联动响应。随着国产芯片生态逐步成熟，类似的端侧智能方案将越来越多地出现在工厂、社区和城市大脑之中。