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

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

一、总体思路

首先在本机上使用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()