rk3399pro_(2)YOLOv5移植

前言

本文介绍yolov5s在瑞芯微rk3399pro上部署推理的过程，并对整个流程以及踩坑记录。

移植流程

torch(.pt)->onnx->rknn

rknn瑞芯微官方文档：github

搭建环境

1. 服务器 训练自己的数据集
   采用U版的yolov5训练，需要安装torch1.9.0+cu111对应版本的依赖包

2. yolov5 训练完成pt->onnx
   安装onnx

3. Linux onnx->rknn （rknn-toolkit-1.7.1）
   此处采用的是1.7.1版本，在上面链接下载官方提供的whl文件，并在python36环境下pip install
   根据瑞芯微提供的文档，转换时需要满足的版本，推荐使用Ubuntu平台
   

4. RK3399pro 移植推理

YOLOv5

此处采用的是U版的yolov5 5.0版本（此时yolov5最新版本已经是）
瑞芯微不同版本对OP的支持不同，所以需要根据网络结构参考版本支持文档

记录：由于之前工作部署hisi平台时，caffe不支持Upsample操作，所以将上采样操作换为反卷积操作。在瑞芯微提供的环境下加载模型时，出现了ConvTranspose reshape fail的错误，查阅提供的支持文档后，发现1.6.0支持反卷积操作，1.7.1似乎不支持，随即又换回了Upsample，加载和模型转换成功。

转换ONNX

yolov5训练完成后，使用yolov5提供的export.py导出onnx格式权重
export.py参数设置：

# data: 加载的模型配置文件
# weights:训练好的权重文件
# opset：ONNX转换参数 设置版本
# simplify：简化ONNX模型转换参数，去除head层和后处理层
# device：使用cpu或gpu转换
# imgsz：导出ONNX参数 输入图像大小
python export.py --data yolov5s.yaml --device 0 --weights best.pt --opset 11 --simplify

转换rknn

在环境搭建中，rknn-tookit-1.7.1安装完成后，在Linux平台下参考example转换

• doc：为参考文档，对load_onnx等加载模型的python Api接口有说明解释
• example：提供一些模型移植转换的参考案例

在example/onnx中，1.7.1是提供了yolov5的转换案例包括后处理代码，此处1.7.1不支持ConvTranspose操作，因此将其作为恢复为Upsample。代码参考test.py

test.py 修改：

ONNX_MODEL = 'best.onnx'
RKNN_MODEL = 'yolov5s.rknn'

IMG_SIZE = 640 # 模型输入大小
CLASSES = (自定义数据类别)
# outputs 为模型三个节点输出的三个尺度，分别是(1,255,80,80)/(1,255,40,40)/(1,255,20,20)
# outputs 参数为节点name，根据netron来查看 
ret = rknn.load_onnx(model=ONNX_MODEL,outputs=['325', '345', '365'])

netron查看ONNX节点输出，以VOC20classname，输出尺寸为40*40的为例

Conv的输出结果即为outputs中需要填写的节点，其name:345，即将outputs=[“xxx”,“345”,“xxx”]，其他两个尺度类似

注意：1.6.0版本的tookit提供的load_onnx函数不支持指定输出层(outputs参数)