官方YOLOV5的torch模型->ONNX模型->RKNN模型

1、环境配置

1.1 RKNN Toolkit2的环境配置

下载RKNN Toolkit2

git clone https://github.com/rockchip-linux/rknn-toolkit2.git

打开一个终端命令行窗口,安装 Python3.6 和 pip3

sudo apt-get install python3 python3-dev python3-pip

安装所需的依赖包

sudo apt-get install libxslt1-dev zlib1g-dev libglib2.0 libsm6 \
libgl1-mesa-glx libprotobuf-dev gcc

进入 Toolkit2 工程文件夹,根据工程的保存路径修改 cd 命令中的路径

cd <输入 Toolkit2 工程的路径>

安装必要相应版本的依赖包

pip3 install -r doc/requirements_cp36-1.x.x.txt

安装 RKNN-Toolkit2(Python3.6 for x86_64)

pip3 install \
package/rknn_toolkit2-1.x.x_xxxxxxxx-cp36-cp36m-linux_x86_64.whl

检查 RKNN-Toolkit2 是否安装成功,在终端输入python后,运行如下python命令

from rknn.api import RKNN

备注:上面的环境配置步骤参考rknn-toolkit2-1.x.0/Rockchip_Quick_Start_RKNN_SDK_V1.4.0_CN.pdf

1.2 Yolov5的环境配置

下载yolov5代码yolov5官方链接,并下载相关的模型,本教程下载的是yolov5s。

创建虚拟环境

conda create -n yolov5 python=3.8

激活虚拟环境

conda activate yolov5

安装python包

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

问题:出现了AttributeError: module ‘numpy’ has no attribute 'typeDict’的错误。
原因分析:numpy的版本太高了。项目中要求安装numpy>=1.18.5,我安装了numpy == 1.24.3版本的。
解决办法:将numpy版本降下来。比如安装1.2.1版本的numpy。把requirements.txt中的numpy>=1.18.5,改为numpy==1.21

2、torch模型->ONNX模型->RKNN模型

2.1 将yolov5s.pt、yolov5m.pt、yolov5x.pt拷贝到下yolov5代码目录官方YOLOV5的torch模型->ONNX模型->RKNN模型_第1张图片

2.2 将torch模型转换成onnx模型

执行以下命令,生成 yolov5s.onnx模型。

python3 export.py --weights yolov5s.pt --img 640 --batch 1 --opset 12 --include onnx

备注:如果在args不加这个–include onnx,会默认生成TorchScript模型

2.3 将onnx模型转换成rknn模型

将2.2生成onnx模型拷贝到./rknn-toolkit2-1.4.0/examples/onnx/yolov5中,我将官方导出onnx模型命名为myyolov5.onnx
官方YOLOV5的torch模型->ONNX模型->RKNN模型_第2张图片修改test.py

# 修改第10行和第11行,指定onnx和rknn模型
ONNX_MODEL = 'myyolov5s.onnx'
RKNN_MODEL = 'myyolov5s.rknn'

# 将第248行ret = rknn.load_onnx(model=ONNX_MODEL)进行修改
ret = rknn.load_onnx(model=ONNX_MODEL,outputs=['/model.24/m.0/Conv_output_0', '/model.24/m.1/Conv_output_0', '/model.24/m.2/Conv_output_0'])

备注:
rknn.load_onnx中的outputs安装输出shape分别是[1, 3, 80, 80, 85],[1, 3, 40, 40, 85],[1, 3, 20, 20, 85]来依次填写节点名字。
具体的网络结构可以通过netron查看

运行test.py

python test.py

官方YOLOV5的torch模型->ONNX模型->RKNN模型_第3张图片

参考链接:Yolov5 转换成 RKNN模型

你可能感兴趣的:(模型量化,YOLO,ONNX,RKNN)