瑞芯微rv1126上实现yolov5模型部署

瑞芯微rv1126上实现yolov5模型部署 保持微笑 2024-02-11 14:06:00 1352

第一次接触模型训练和在开发板部署,过程曲折,从开始的一脸懵到最后模型部署成功,查阅了不少资料和学习了不少大佬的经验,在这里记录一下过程和心得。

一、总体思路

首先在本机上使用yolov5代码训练自己的数据集,得到训练后的pt模型,然后通过export.py转化出onnx模型,最后通过瑞芯微官方代码,将onnx模型转化为rknn模型,实现在rv1126开发板上实现模型部署。

ps:为什么不直接使用pytorch模型转化为rknn?
答:rknn1.7.1直接加载pt模型会有bug。

二、实现过程

1、yolov5训练

我是在本机windows训练的模型,所以环境也是装在windows上,创建环境我用的是Anaconda,具体怎么下载安装有很多成功的博客,麻烦大家自己搜索。

首先拉取官方代码,我使用的yolov5代码版本是v5.0 release版本,项目地址是:https://github.com/ultralytics/yolov5/tree/v5.0

  • conda创建环境指令如下:(python版本是3.7)
conda create --name yolov5 python==3.7
  • 配置训练所需的环境
进入虚拟环境:
activate yolov5

推荐使用清华源下载:
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt
  • 数据标注准备

    • 安装labelimg
      pip install labelimg
    • 使用labelimg标注
      点击Open Dir -> 选择需要标注的文件夹 → ok
      进行标注格式切换,切换到yolo。
      点击Create RectBox → 拖拽鼠标框选目标 → 给上标签 → 点击ok。
      点击save,保存txt。
  • 修改yolov5项目文件
    1、yolov5/data文件夹下新建fire.yaml(名字自取),用来存放训练集和验证集的划分文件(train.txt和val.txt),参考如下:

    train: D:\Lenovo\Desktop\v5-5.0\paper_date\images\train\   
    val: D:\Lenovo\Desktop\v5-5.0\paper_date\images\val\
    # Classes   
    nc: 1  # number of classes   
    names: ['plasticbottle']  # class names
    

    2、model下修改yolov5s.yaml的种类数nc=1,(因为我这里只训练了一种类别,具体根据自己需要改)

3、如果显卡显示不足的可以将dataset.py中的num_worker数量设置为0.

4、yolo.py修改detect设置nc的值,并且替换掉dectet模块
参考:https://zhuanlan.zhihu.com/p/450411127 第三点

5、common.py补充缺少的模块定义
参考:https://zhuanlan.zhihu.com/p/450411127 第二点

6、export.py添加参数 parser.add_argument(‘–grid’, action=‘store_true’, help=‘export Detect() layer grid’)

7、(*)这点很重要,因为rknn官方代码不支持onnx输出的Detect层,因此在运行完train.py训练出pt模型后,使用export.py前需要将yolo.py中的代码

x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
修改为
x[i] = x[i].view(bs, self.na, self.no, ny, nx)

而当转化出onnx模型后,若想重新训练新的模型,需要将对yolo.py上述改动修改回去,否则会出现报错!!

参考:35、ubuntu20.04搭建瑞芯微的npu仿真环境和测试rv1126的Debain系统下的yolov5+npu检测功能以及RKNN推理部署

train.py参数设置:

上述做完就可以开始愉快的训练和转onnx模型了。
以下是我遇到的报错:
报错1

OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized.

在pycharm里调试程序时可以直接通过在程序前添加这两个语句解决

import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

报错2

RuntimeError: result type Float can't be cast to the desired output type __int64

参考:https://blog.csdn.net/Thebest_jack/article/details/125649451

报错3

RuntimeError: The size of tensor a (52) must match the size of tensor b (84) at non-singleton dimension 3

 参考:https://blog.csdn.net/Thebest_jack/article/details/125649451
 (这就是我上面说的训练没有将yolo.py中onnx的改动修改回去发生的报错。)

2、模型转化(onnx–>rknn)

我这里是将rknn官方的代码拆分成两个部分,一是模型转化,二就是在板子上进行后处理以及推理。
这里贴出onnx转化rknn的代码(通过官方代码修改),量化需要的数据集和标签自己修改。

import os
import urllib
import traceback
import time
import sys
from rknn.api import RKNN
ONNX_MODEL = 'best.onnx'
RKNN_MODEL = 'yolov5s.rknn'
DATASET = './dataset.txt'

QUANTIZE_ON = True

BOX_THRESH = 0.5
NMS_THRESH = 0.6
IMG_SIZE = 640

if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN()

    if not os.path.exists(ONNX_MODEL):
        print('model not exist')
        exit(-1)

    # pre-process config
    print('--> Config model')
    rknn.config(reorder_channel='0 1 2',
                mean_values=[[0, 0, 0]],
                std_values=[[255, 255, 255]],
                optimization_level=3,
                target_platform='rv1126',
                output_optimize=1,
                quantize_input_node=QUANTIZE_ON)
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,outputs=['output', '405', '418']) 
    if ret != 0:
        print('Load yolov5 failed!')
        exit(ret)
    print('done')

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

    # Export RKNN model
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export yolov5rknn failed!')
        exit(ret)
    print('done')

    rknn.release()

注:先使用Netron查看已经转化好的onnx模型结构,看到OUTPUTS的值。从而修改上述代码中rknn.load_onnx的参数。如下图:

环境的配置有需要的可以参考这篇博客:RKNN-toolkit-1.7.1 安装踩坑记录(安装成功)

3、部署

将转化好的rknn模型上传到开发板上,官方的代码有点bug,可以参考一个大佬的博客:35、ubuntu20.04搭建瑞芯微的npu仿真环境和测试rv1126的Debain系统下的yolov5+npu检测功能以及RKNN推理部署,在第七步左右。

到这里部署就基本完成了,只是在第一次在板子上检测图片的时候发现预测结果没有框,后来发现是给的训练数据集太少了,加上模型经过量化,精度下降,所以没有得出理想结果。
解决办法:扩大数据集的规模,用更多的图片进行训练,提高精度。

最后附上一张检测图片:

声明:本文内容由易百纳平台入驻作者撰写,文章观点仅代表作者本人,不代表易百纳立场。如有内容侵权或者其他问题,请联系本站进行删除。
红包 点赞 收藏 评论 打赏
评论
0个
内容存在敏感词
手气红包
    易百纳技术社区暂无数据
相关专栏
置顶时间设置
结束时间
删除原因
  • 广告/SPAM
  • 恶意灌水
  • 违规内容
  • 文不对题
  • 重复发帖
打赏作者
易百纳技术社区
保持微笑
您的支持将鼓励我继续创作!
打赏金额:
¥1易百纳技术社区
¥5易百纳技术社区
¥10易百纳技术社区
¥50易百纳技术社区
¥100易百纳技术社区
支付方式:
微信支付
支付宝支付
易百纳技术社区微信支付
易百纳技术社区
打赏成功!

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

举报反馈

举报类型

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

详细说明

审核成功

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

审核失败

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

小包子的红包

恭喜发财,大吉大利

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

    易百纳技术社区