前言
本文介绍yolov5s在瑞芯微rk3399pro上部署推理的过程,并对整个流程以及踩坑记录。
移植流程
torch(.pt)->onnx->rknn
rknn瑞芯微官方文档:github
搭建环境
-
服务器 训练自己的数据集
采用U版的yolov5训练,需要安装torch1.9.0+cu111对应版本的依赖包 -
yolov5 训练完成pt->onnx
安装onnx -
Linux onnx->rknn (rknn-toolkit-1.7.1)
此处采用的是1.7.1版本,在上面链接下载官方提供的whl文件,并在python36环境下pip install
根据瑞芯微提供的文档,转换时需要满足的版本,推荐使用Ubuntu平台
-
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参数)