从YOLOv5到RKNN:零冲突转换YOLOv5模型至RK3588 NPU全指南

从YOLOv5到RKNN:零冲突转换YOLOv5模型至RK3588 NPU全指南

在嵌入式AI领域,将训练好的深度学习模型高效部署到边缘设备的NPU(神经网络处理器)上是提升性能的关键。本文将详细介绍如何在Ubuntu 20.04环境下,将YOLOv5l模型(.pt文件)通过ONNX中间格式,最终转换为适用于RK3588 NPU的RKNN模型(.rknn文件),并全程采用环境隔离技术确保无依赖冲突。

一、准备工作:构建干净的工作环境

环境污染是模型转换过程中最常见的问题之一,因此我们首先创建一个全新的工作区,确保所有操作在干净的环境中进行。

1. 清理与创建工作区

打开终端,执行以下命令清理可能存在的旧工作区并创建新的工作目录:

bash 复制代码
# 回到主目录
cd ~
# 删除旧的工作区(如果存在)
rm -rf rknn_workspace
# 创建新工作区并进入
mkdir rknn_workspace
cd rknn_workspace

2. 准备核心文件

在新工作区内,我们需要准备四个关键组成部分:YOLOv5源代码、预训练模型、校准图片集和图片列表文件。

bash 复制代码
# 1. 克隆YOLOv5官方仓库
git clone https://github.com/ultralytics/yolov5.git

# 2. 下载yolov5l预训练模型
wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l.pt

# 3. 创建图片文件夹(用于模型量化校准)
mkdir images

# 4. 生成图片列表文件(后续量化会用到)
find "$(pwd)/images" -name "*.jpg" > dataset.txt
find "$(pwd)/images" -name "*.png" >> dataset.txt

完成后,执行ls命令检查,确保工作区内包含以下内容:

  • yolov5(文件夹)
  • yolov5l.pt(模型文件)
  • images(文件夹)
  • dataset.txt(文件)

二、任务一:导出ONNX模型(隔离环境操作)

ONNX(Open Neural Network Exchange)是一种通用的模型中间格式,我们首先将PyTorch模型转换为ONNX格式。为避免依赖冲突,这一步将在独立的虚拟环境中进行。

1. 创建并激活虚拟环境

bash 复制代码
# 进入YOLOv5目录
cd yolov5
# 创建虚拟环境
python3 -m venv onnx_export_venv
# 激活虚拟环境
source onnx_export_venv/bin/activate

激活成功后,终端提示符前会显示(onnx_export_venv),表示当前处于隔离环境中。

2. 安装依赖库

为加快下载速度并确保版本兼容,我们使用清华源并锁定特定版本的ONNX库:

bash 复制代码
# 配置清华源
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
# 升级pip
pip install --upgrade pip
# 安装YOLOv5依赖
pip install -r requirements.txt
# 安装特定版本的ONNX(与RKNN工具兼容)
pip install onnx==1.12.0

3. 导出ONNX模型

将预训练模型复制到当前目录并执行导出命令:

bash 复制代码
# 复制模型文件到当前目录
cp ../yolov5l.pt .

# 导出ONNX模型(指定opset 11以确保兼容性)
python export.py --weights yolov5l.pt --include onnx --opset 11

4. 整理结果并清理环境

导出成功后,会生成yolov5l.onnx文件。我们将其移至工作区根目录,然后销毁当前虚拟环境(使命已完成):

bash 复制代码
# 将ONNX模型移至工作区根目录
mv yolov5l.onnx ../
# 退出虚拟环境
deactivate
# 返回工作区根目录
cd ..
# 删除虚拟环境(释放空间)
rm -rf ./yolov5/onnx_export_venv

此时,工作区根目录应已出现yolov5l.onnx文件,第一步转换完成。

三、任务二:转换为RKNN模型(独立隔离环境)

接下来,我们在全新的虚拟环境中使用RKNN-Toolkit2将ONNX模型转换为RK3588可识别的RKNN模型。

1. 创建并激活新虚拟环境

bash 复制代码
# 确保在rknn_workspace目录下
python3 -m venv rknn_convert_venv
# 激活环境
source rknn_convert_venv/bin/activate

终端提示符前会显示(rknn_convert_venv),表示进入新的隔离环境。

2. 安装RKNN-Toolkit2

bash 复制代码
# 配置清华源
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
# 升级pip
pip install --upgrade pip
# 安装指定版本的RKNN-Toolkit2(与RK3588兼容)
pip install rknn-toolkit2==2.3.2

该命令会自动安装所有依赖的兼容版本库,无需手动干预。

3. 准备转换脚本

创建convert_rknn.py脚本,用于执行模型转换:

bash 复制代码
nano convert_rknn.py

复制以下代码到文件中:

python 复制代码
from rknn.api import RKNN

if __name__ == '__main__':
    rknn = RKNN(verbose=True)

    print('--> Config model')
    rknn.config(
        mean_values=[[0, 0, 0]],
        std_values=[[255, 255, 255]],
        target_platform='rk3588'  # 指定目标平台为RK3588
    )
    print('done')

    print('--> Loading model')
    ret = rknn.load_onnx(model='./yolov5l.onnx')
    if ret != 0:
        print('!!! Load yolov5l.onnx failed!')
        exit(ret)
    print('done')

    print('--> Building model')
    # 启用量化,并指定校准数据集
    ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
    if ret != 0:
        print('!!! Build yolov5l.rknn failed!')
        exit(ret)
    print('done')

    print('--> Export rknn model')
    ret = rknn.export_rknn('./yolov5l.rknn')
    if ret != 0:
        print('!!! Export yolov5l.rknn failed!')
        exit(ret)
    print('done')

    rknn.release()

保存并退出(按Ctrl+O,回车,再按Ctrl+X)。

4. 准备量化校准图片

模型量化是提升NPU运行效率的关键步骤,需要20-30张代表性图片作为校准数据。如果没有现成图片,可以从视频中提取:

(1)安装OpenCV用于视频帧提取
bash 复制代码
pip install opencv-python
(2)创建视频帧提取脚本
bash 复制代码
nano extract_frames.py

复制以下代码:

python 复制代码
import cv2
import os
import sys

# 配置区 - 修改为你的视频路径
VIDEO_PATH = "/path/to/your/video.mp4"  # 替换为实际视频路径
OUTPUT_DIR = "images"
NUM_FRAMES_TO_EXTRACT = 30  # 提取30帧

def extract_frames():
    if not os.path.exists(VIDEO_PATH):
        print(f"错误: 视频文件未找到! 路径: {VIDEO_PATH}")
        sys.exit(1)

    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)

    cap = cv2.VideoCapture(VIDEO_PATH)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    if total_frames == 0:
        print("错误: 无法读取视频或视频为空。")
        cap.release()
        sys.exit(1)

    # 计算采样间隔,确保均匀取帧
    interval = total_frames // NUM_FRAMES_TO_EXTRACT
    if interval == 0:
        interval = 1

    print(f"视频总帧数: {total_frames}")
    print(f"计划提取 {NUM_FRAMES_TO_EXTRACT} 帧, 间隔 {interval} 帧。")

    frame_count = 0
    extracted_count = 0
    while cap.isOpened() and extracted_count < NUM_FRAMES_TO_EXTRACT:
        ret, frame = cap.read()
        if not ret:
            break

        if frame_count % interval == 0:
            output_filename = os.path.join(OUTPUT_DIR, f"frame_{extracted_count:04d}.jpg")
            cv2.imwrite(output_filename, frame)
            print(f"已保存: {output_filename}")
            extracted_count += 1

        frame_count += 1

    cap.release()
    print(f"\n提取完成! 共保存 {extracted_count} 张图片到 '{OUTPUT_DIR}'。")

if __name__ == "__main__":
    extract_frames()
(3)执行帧提取并更新数据集列表
bash 复制代码
# 修改脚本中的VIDEO_PATH为实际视频路径后执行
python3 extract_frames.py

# 更新dataset.txt文件
find "$(pwd)/images" -name "*.jpg" > dataset.txt
find "$(pwd)/images" -name "*.png" >> dataset.txt

5. 执行RKNN转换

一切准备就绪后,运行转换脚本:

bash 复制代码
python3 convert_rknn.py

转换过程会持续几分钟,成功完成后会显示类似以下的日志:

复制代码
I rknn-toolkit2 version: 2.3.2
...
I rknn building done.
done
--> Export rknn model
done

此时,工作区根目录会生成yolov5l.rknn文件,这就是RK3588 NPU可直接使用的模型文件。

四、常见问题与解决方案

  1. 错误:Dataset file ./dataset.txt not found!

    • 原因:未生成数据集列表文件或文件路径错误
    • 解决:确保images文件夹中有图片,并重新执行dataset.txt生成命令
  2. ONNX导出失败

    • 原因:YOLOv5版本与依赖库不兼容
    • 解决:使用本文指定的YOLOv5仓库和onnx==1.12.0版本
  3. RKNN转换时量化失败

    • 原因:校准图片数量不足或与模型场景不匹配
    • 解决:确保提供20-30张与模型应用场景相符的图片(如检测交通场景则提供交通图片)

五、总结与后续步骤

通过本文介绍的两步隔离环境转换法,我们成功将YOLOv5l模型转换为RK3588 NPU适用的RKNN模型,全程避免了依赖冲突问题。转换后的yolov5l.rknn文件可直接部署到RK3588设备。

后续步骤建议:

  1. yolov5l.rknn传输到已安装Ubuntu系统的RK3588开发板
  2. 使用RKNN-Toolkit2或RKNN Runtime API编写推理脚本
  3. 在RK3588上测试模型性能与精度,必要时进行优化调整

这种隔离环境的转换方法不仅适用于YOLOv5,也可推广到其他深度学习模型的RKNN转换过程,为嵌入式AI部署提供了可靠的技术参考。

相关推荐
Sadsvit1 小时前
网络聚合链路与软件网桥配置指南
linux·运维·服务器·网络·centos
___波子 Pro Max.1 小时前
Linux数字列排序命令
linux
Kusunoki_D1 小时前
PyTorch 环境配置
人工智能·pytorch·python
飞凌嵌入式2 小时前
高性能、高实时、高安全:如何在飞凌嵌入式i.MX95xx核心板上同时实现?
嵌入式硬件·安全·嵌入式·飞凌嵌入式
知秋丶2 小时前
大模型应用发展与Agent前沿技术趋势(下)
人工智能·python·ai agent
--运维实习生--3 小时前
shell脚本第二阶段-----选择结构
linux·运维·shell编程
Clownseven3 小时前
Linux服务器健康检查Shell脚本:一键生成自动化巡检报告
linux·服务器·自动化
智能物联实验室3 小时前
如何低门槛自制Zigbee 3.0温湿度计?涂鸦上新开发包,开箱即用、完全开源
嵌入式硬件·开源·硬件工程
HenryLin3 小时前
美股量化分析系统 - 模块调用流程文档
python
跟橙姐学代码4 小时前
一文读懂 Python 的 JSON 模块:从零到高手的进阶之路
前端·python