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

在智能安防、工业质检和边缘视觉日益普及的今天，如何将高效的目标检测模型稳定落地到嵌入式设备上，已成为AI工程化的核心挑战。瑞芯微RK3588凭借其6TOPS INT8算力的NPU和四核A76+四核A55的异构架构，正成为国产边缘计算平台中的“性能担当”。而YOLOv5以其轻量级设计和出色的精度-速度平衡，自然成为开发者首选。

但理论归理论，真正从数据准备到NPU推理跑通全流程，中间仍有不少“坑”——比如SiLU激活函数在NPU上的兼容性问题、ONNX导出失败、量化精度下降等。本文不讲空泛概念，而是以一次真实项目实践为蓝本，带你完整走一遍基于RK3588的自定义目标检测模型训练与部署全过程，涵盖环境搭建、数据处理、模型调优、格式转换及实机推理优化等关键环节。

---

我们从一个实际场景切入：某智能洗手台项目需要识别用户是否伸手（hand）、露脸（face）或手持水杯（cup），以便触发自动出水/语音提示功能。这正是典型的多类别小样本目标检测任务，非常适合用YOLOv5 + RK3588组合来实现。

首先面临的第一个问题是：开发环境怎么搭？

Python生态虽然强大，但深度学习项目依赖复杂，不同框架版本之间极易冲突。直接用系统Python很容易“污染”全局环境。因此我们采用 Miniconda-Python3.10镜像 构建隔离环境。它体积小、启动快，且能精确控制每个项目的包版本，特别适合需要复现结果的AI研发。

如果你习惯交互式调试，可以安装Jupyter：

conda install jupyter -y
jupyter notebook --ip=0.0.0.0 --port=8888 --allow-root --no-browser

通过浏览器即可访问Notebook界面，适合做数据探索和模型验证。但对于长期运行的训练任务，更推荐使用SSH登录远程服务器或开发主机进行命令行操作，资源管理更清晰也更安全。

接下来创建专属虚拟环境：

conda create -n rk3588_yolov5 python=3.10 -y
conda activate rk3588_yolov5

环境准备好后，重点来了：必须使用适配RK3588 NPU的YOLOv5分支！

GitHub地址：https://github.com/airockchip/yolov5

为什么不能用官方ultralytics/yolov5？因为后者默认使用SiLU（Swish）作为激活函数，而RK3588的NPU对ReLU支持更好，SiLU可能导致算子不支持或推理性能大幅下降。这个细节看似微小，实则决定了整个项目能否顺利部署。

克隆代码并安装依赖：

git clone https://github.com/airockchip/yolov5.git
cd yolov5
pip install -r requirements.txt

如果有GPU用于加速训练，建议安装CUDA版PyTorch：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

验证一下环境是否正常：

import torch
print(torch.cuda.is_available())  # 应输出 True
print(torch.__version__)

再执行 nvidia-smi 查看显卡状态，确认驱动和CUDA都已就绪。

---

接下来是数据准备阶段。没有高质量的数据，再好的模型也是空中楼阁。

我们的目标是构建一个包含“手”、“脸”、“杯子”三类对象的小型数据集。图像采集时要注意覆盖多种光照条件（白天/夜晚/背光）、不同角度（俯视/平视）以及部分遮挡情况，这样才能提升模型泛化能力。

推荐目录结构如下：

dataset/
├── images/
│   └── *.jpg
└── annotations/
    └── *.xml

标注工具推荐使用 LabelImg，它是开源的图形化标注软件，支持PASCAL VOC格式输出（即XML文件）。安装方式简单：

pip install labelme

启动后操作流程也很直观：
1. 打开图像目录
2. 设置保存路径
3. 切换至 PascalVOC 模式
4. 绘制边界框并输入类别名（如 hand, face, cup）

生成的XML示例如下：

<annotation>
    <folder>images</folder>
    <filename>img_001.jpg</filename>
    <path>/home/user/dataset/images/img_001.jpg</path>
    <size>
        <width>640</width>
        <height>480</height>
        <depth>3</depth>
    </size>
    <object>
        <name>hand</name>
        <bndbox>
            <xmin>100</xmin>
            <ymin>80</ymin>
            <xmax>250</xmax>
            <ymax>230</ymax>
        </bndbox>
    </object>
</annotation>

但YOLO系列模型要求标签为.txt格式，内容为归一化的中心坐标和宽高（class_id x_center y_center width height）。为此我们需要编写一个转换脚本，同时完成数据集划分：

import os
import xml.etree.ElementTree as ET
import random
import shutil

classes = ['hand', 'face', 'cup']
TRAIN_RATIO = 80

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_name):
    in_path = f'./dataset/annotations/{image_name}.xml'
    out_path = f'./dataset/labels/{image_name}.txt'

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

    with open(out_path, 'w') as out_file:
        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(f"{cls_id} {' '.join([str(round(x, 6)) for x in bb])}\n")

os.makedirs('./dataset/labels/train', exist_ok=True)
os.makedirs('./dataset/labels/val', exist_ok=True)
os.makedirs('./dataset/images/train', exist_ok=True)
os.makedirs('./dataset/images/val', exist_ok=True)

image_files = [f[:-4] for f in os.listdir('./dataset/annotations') if f.endswith('.xml')]
random.shuffle(image_files)

with open('train.txt', 'w') as tf, open('val.txt', 'w') as vf:
    for i, image_name in enumerate(image_files):
        img_src = f'./dataset/images/{image_name}.jpg'
        if not os.path.exists(img_src):
            continue
        convert_annotation(image_name)

        dst_dir = 'train' if i < len(image_files) * TRAIN_RATIO / 100 else 'val'
        shutil.copy(img_src, f'./dataset/images/{dst_dir}/')
        shutil.copy(f'./dataset/labels/{image_name}.txt', f'./dataset/labels/{dst_dir}/')
        (tf if dst_dir == 'train' else vf).write(img_src + '\n')

执行后会自动生成 train.txt 和 val.txt 文件，并按比例复制图像和标签到对应子目录，省去手动整理的麻烦。

---

进入模型训练阶段前，先配置 data/custom.yaml：

train: ./train.txt
val: ./val.txt
nc: 3
names: ['hand', 'face', 'cup']

关于模型尺寸选择，虽然YOLOv5提供了s/m/l/x多个版本，但在嵌入式场景下强烈建议使用 yolov5s。它的参数量仅7.2M，推理速度快，且在RK3588上实测可达25+ FPS（INT8量化后），完全满足实时性需求。更大的模型不仅加载慢，还可能因内存不足导致崩溃。

开始训练命令如下：

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

其中 --cache 参数会将图像缓存到内存中，避免重复IO读取，可显著加快训练速度，尤其适合小数据集。如果显存足够，也可以适当增大batch size以提高梯度稳定性。

训练过程中可通过TensorBoard监控loss曲线和mAP变化，判断是否过拟合。一般50~100轮即可收敛。最终最优模型会保存在 runs/train/yolov5s_rk3588/weights/best.pt。

---

模型训练完成后，下一步是将其部署到RK3588设备上。由于NPU只能运行专有格式的模型，我们必须经过两步转换：ONNX导出 → RKNN量化。

首先导出ONNX模型：

python export.py --weights runs/train/yolov5s_rk3588/weights/best.pt --include onnx

生成的 best.onnx 可用Netron打开查看网络结构，确认输入输出节点无误。

然后安装PC端的RKNN转换工具链：

pip install rknn-toolkit-lite2

编写转换脚本 onnx_to_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)

ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
if ret != 0:
    print('Build RKNN failed!')
    exit(ret)

rknn.export_rknn('yolov5s_custom.rknn')
rknn.release()

这里的关键是 dataset.txt —— 它需要包含至少50张真实图像路径（每行一条），用于校准量化过程。太少会导致精度损失严重，太多则耗时增加。建议从中随机采样一批训练集图像生成该文件。

转换成功后，将 yolov5s_custom.rknn 推送到RK3588开发板：

scp yolov5s_custom.rknn root@192.168.1.100:/userdata/

在设备端安装运行时库：

pip install rknn-runtime

编写推理脚本 infer.py：

import cv2
import numpy as np
from rknn.api import RKNN

rknn = RKNN()

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

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

img = cv2.imread('test.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (640, 640))
img = np.expand_dims(img, axis=0)

outputs = rknn.inference(inputs=[img])
print("Output shapes:", [o.shape for o in outputs])

rknn.release()

注意几个细节：
- 图像需转为RGB顺序（OpenCV默认BGR）
- 输入要扩展batch维度
- core_mask=RKNN.NPU_CORE_0_1_2 启用三核并行，可提升约2.5倍吞吐

推理输出是一个包含三个tensor的列表，分别对应三个检测头的特征图，后续可通过非极大抑制（NMS）解析出最终检测框。

---

整个流程走完你会发现，真正的难点往往不在算法本身，而在软硬件协同的细节把控。比如一次实际测试中，我们发现INT8量化后mAP掉了近8个百分点，排查后才发现是预处理均值/方差设置错误——应该用 [0,0,0] 和 [255,255,255] 进行归一化，而不是常见的ImageNet统计值。

另一个经验是：对于小目标较多的场景（如指尖检测），可以把输入分辨率从640提升到896甚至1280，虽然推理稍慢，但召回率明显改善。RK3588的NPU对大尺寸输入仍有不错的处理能力。

此外，还可以结合OpenCV实现视频流实时检测：

cap = cv2.VideoCapture(0)
while True:
    ret, frame = cap.read()
    if not ret:
        break
    input_img = preprocess(frame)  # resize, normalize, expand dims
    results = rknn.inference(inputs=[input_img])
    boxes = postprocess(results)  # decode & NMS
    draw_boxes(frame, boxes)
    cv2.imshow('RK3588 YOLOv5', frame)
    if cv2.waitKey(1) == ord('q'):
        break

这种低延迟、本地化的视觉处理能力，正是边缘AI的核心价值所在。

---

回顾整个流程，从环境隔离到模型选型，从数据规范到量化校准，每一步都有其技术深意。RK3588 + YOLOv5的组合不仅性能强劲，而且生态成熟，非常适合快速原型开发和产品落地。未来还可进一步尝试模型蒸馏、动态输入、多模型串联等高级优化手段，持续挖掘硬件潜力。

当你看到开发板屏幕上流畅地框出每一个“手”、“脸”、“杯子”，那种从代码到现实的连接感，或许才是做AI最令人着迷的部分。