从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部署提供了可靠的技术参考。

相关推荐
MAR-Sky12 分钟前
单片机学习中的一些简单总结
单片机·嵌入式硬件·学习
点灯小铭14 分钟前
基于stm32的物联网OneNet火灾报警系统
stm32·单片机·嵌入式硬件·物联网·毕业设计·课程设计
hui函数27 分钟前
python全栈(基础篇)——day04:后端内容(字符编码+list与tuple+条件判断+实战演示+每日一题)
开发语言·数据结构·python·全栈
sheji34166 小时前
【开题答辩全过程】以 python杭州亚运会数据分析与可视化开题为例,包含答辩的问题和答案
开发语言·python·数据分析
古月-一个C++方向的小白7 小时前
Linux——查看与创建进程
linux·运维·服务器
许商8 小时前
【stm32】【edgetx】解析链接脚本文件(ld)
stm32·单片机·嵌入式硬件
vortex58 小时前
fd 工具指南:find 的现代替代品
linux·运维开发·命令行工具
2401_841495649 小时前
【计算机视觉】基于数学形态学的保留边缘图像去噪
人工智能·python·算法·计算机视觉·图像去噪·数学形态学·边缘保留
馨谙9 小时前
vim保姆级使用,操作详解,快捷键大全总结
linux·编辑器·vim
驱动探索者9 小时前
find 命令使用介绍
java·linux·运维·服务器·前端·学习·microsoft