yolov5 onnx模型 转为 rknn模型

yolov5 onnx模型 转为 rknn模型 技术小宅 2024-01-06 09:02:48 444

onnx模型需要转换为rknn模型才能在rv1126开发板上运行,所以需要先搭建转换环境

1、转换为rknn模型环境搭建

模型转换工具 模型转换相关文件下载:
网盘下载链接:百度网盘 请输入提取码 提取码:teuc

将其移动到虚拟机中,找到文件夹中的docker文件 rknn-toolkit-1.7.1-docker.tar.gz 、 model_convert文件夹

加载 模型转换工具docker镜像

docker load --input /home/developer/rknn-toolkit/rknn-toolkit-1.7.1-docker.tar.gz

进入 镜像bash环境
执行以下指令把工作区域映射进docker镜像,其中/home/developer/rknn-toolkit/model_convert为工作区域,/test为映射到docker镜像,/dev/bus/usb:/dev/bus/usb为映射usb到docker镜像:

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb -v /home/developer/rknn-toolkit/model_convert:/test rknn-toolkit:1.7.1 /bin/bash

2、生成量化图片列表

这一步是把准备好的一些图片,生成图片路径的文本文件,在构建RKNN模型的时候有用处。通过使用真实的样本数据集,RKNN工具可以更好地理解和建模模型的输入数据,从而更好地优化网络结构、权重和量化方案。

在docker环境切换到模型转换工作目录;执行 gen_list.py,会得到一个文本文件pic_path.txt,里面是图片的路径:

cd /test/coco_object_detect
python gen_list.py

gen_list.py内容如下:

import os
import random


def main(image_dir):
    save_image_txt = './pic_path.txt'
    save_val_number = 0

    img_path_list = []

    image_list = os.listdir(image_dir)
    for i in image_list:
        #if os.path.isdir(image_dir):
        #print("i:", i)
        image_path = image_dir + '/' + i
        #print("image_path:", image_path)
        img_path_list.append(image_path)

    #print(img_path_list)


    print('len of all', len(img_path_list))

    random.shuffle(img_path_list)

    with open(save_image_txt, 'w') as F:
        for i in range(len(img_path_list)):
            F.write(img_path_list[i]+'\n')


if __name__ == '__main__':
    image_dir = '/test/quant_dataset/coco_data'  # 图片所在路径,大概500main(image_dir)

3、onnx模型转换为rknn模型

还是在docker环境 模型转换工作目录,运行rknn_convert.py

python rknn_convert.py

这一步如果是在虚拟机上运行的话,8GB的内存条win10系统也要用,分配给虚拟机的没多少,3GB也不够执行这一步。

后来我直接在Ubuntu系统执行这一步,8GB系统用一点还剩6.7GB,CPU和内存直接干满

rknn_convert.py 源码:

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN


ONNX_MODEL = 'best.onnx' # onnx 模型的路径
RKNN_MODEL = './yolov5_mask_rv1126.rknn'  # 转换后的 RKNN 模型保存路径
DATASET = './pic_path.txt'   # 数据集文件路径

QUANTIZE_ON = True   # 是否进行量化

if __name__ == '__main__':

    # 创建 RKNN 对象
    rknn = RKNN(verbose=True)

    # 检查 ONNX 模型文件是否存在
    if not os.path.exists(ONNX_MODEL):
        print('model not exist')
        exit(-1)

    # 配置模型预处理参数
    print('--> Config model')
    rknn.config(reorder_channel='0 1 2', # 表示 RGB 通道
                mean_values=[[0, 0, 0]], # 每个通道的像素均值,预处理时对应通道减去该值
                std_values=[[255, 255, 255]], # 每个通道的像素标准差,每个通道除以该值
                optimization_level=3, # 优化级别
                target_platform = 'rv1126', #指定目标平台为rv1126
                output_optimize=1,      # 输出优化为真
                quantize_input_node=QUANTIZE_ON)  # 对时输入节点进行量化
    print('done')

    # 加载 ONNX 模型
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load yolov5 failed!')
        exit(ret)
    print('done')

    # 构建模型
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build yolov5 failed!')
        exit(ret)
    print('done')

    # 导出 RKNN 模型
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export yolov5rknn failed!')
        exit(ret)
    print('done')
声明:本文内容由易百纳平台入驻作者撰写,文章观点仅代表作者本人,不代表易百纳立场。如有内容侵权或者其他问题,请联系本站进行删除。
红包 2 1 评论 打赏
评论
0个
内容存在敏感词
手气红包
    易百纳技术社区暂无数据
相关专栏
置顶时间设置
结束时间
删除原因
  • 广告/SPAM
  • 恶意灌水
  • 违规内容
  • 文不对题
  • 重复发帖
打赏作者
易百纳技术社区
技术小宅
您的支持将鼓励我继续创作!
打赏金额:
¥1易百纳技术社区
¥5易百纳技术社区
¥10易百纳技术社区
¥50易百纳技术社区
¥100易百纳技术社区
支付方式:
微信支付
支付宝支付
易百纳技术社区微信支付
易百纳技术社区
打赏成功!

感谢您的打赏,如若您也想被打赏,可前往 发表专栏 哦~

举报反馈

举报类型

  • 内容涉黄/赌/毒
  • 内容侵权/抄袭
  • 政治相关
  • 涉嫌广告
  • 侮辱谩骂
  • 其他

详细说明

审核成功

发布时间设置
发布时间:
是否关联周任务-专栏模块

审核失败

失败原因
备注
拼手气红包 红包规则
祝福语
恭喜发财,大吉大利!
红包金额
红包最小金额不能低于5元
红包数量
红包数量范围10~50个
余额支付
当前余额:
可前往问答、专栏板块获取收益 去获取
取 消 确 定

小包子的红包

恭喜发财,大吉大利

已领取20/40,共1.6元 红包规则

    易百纳技术社区