yolov5-6.2 在 rk3399pro 上的移植
文章目录
- 一、搭建yolov5环境
- 二、导出onnx模型
- 三、安装 rknn-toolkit
- 四、通过netron 查看 yolov5s.onnx 输出节点
- 五、创建 onnx2rknn.py 文件
- 六、通过toolkit将onnx转换为rknn模型
- 七、在rk3399pro开发板上运行rknn模型
一、搭建yolov5环境
从yolov5官方仓库 ultralytics/yolov5/releases 中下载 yolov5-6.2 代码。
解压后，修改 requirements.txt :
1、将第12行 torch>=1.7.0 改为 torch==1.9.0
因为只有 pytorch 1.8.0 或 1.9.0 版本才能正常导出模型
2、删除第28、29行 onnx>=1.9.0 前的 # 号，否则无法导出onnx模型；
执行
pip install -r requirement.txt
验证环境是否搭建成功
python detect.py
二、导出onnx模型
1、进入 yolov5-6.2/models ，将 yolo.py 文件中 class Detect 中的 forward 函数，修改如下
def forward(self, x):z = [] # inference outputfor i in range(self.nl):x[i] = self.m[i](x[i]) # convreturn x[0],x[1],x[2]
2、进入 yolov5-6.2 主目录，执行以下命令导出onnx模型
python export.py --weights yolov5s.pt --img 640 --batch 1 --opset 12 --include onnx
执行成功后，将在 yolov5s.pt 同目录生成 yolov5s.onnx 文件
三、安装 rknn-toolkit
1、下载docker版本的toolkit
从 rockchip-linux/rknn-toolkit 下载 docker版本的toolkit，提取码：rknn
这里下载的是 1.7.3 版本
2、在安装有docker的linux服务器上导入上述toolkit
docker load --input rknn-toolkit-1.7.3-docker.tar.gz
导入完成后执行 docker images，将看到 toolkit-1.7.3 镜像
执行以下命令，进入docker
docker run -t -i --privileged -v $PWD:/test rknn-toolkit:1.7.3 /bin/bash
其中 -v $PWD:/test 是指把当前目录映射到docker容器的/test目录内，这样当要和docker容器内部进行文件传输时，只需把需要传输的文件放到当前目录，就能在 docker 容器内的 /test 目录看到，反之亦然。
执行成功后，将自动进入docker容器内
可以看到已经有一个 test 目录，test 目录的内容就是服务器当前目录的内容。
四、通过netron 查看 yolov5s.onnx 输出节点
先安装后 netron，然后netron打开上述第二步导出的 yolov5s.onnx 模型，点击左上角图标，即可看到输出节点信息
可以看到这里的输出节点分别是: output、327、328，记住这几个输出节点，下一步导出rknn时将用到
五、创建 onnx2rknn.py 文件
在该文件中指定onnx_model路径，platform，以及 rknn.load_onnx 时一定要指定上一步查看的输出节点信息
from rknn.api import RKNN
ONNX_MODEL = 'yolov5s.onnx'
platform = "rk3399pro"
RKNN_MODEL = 'yolov5s_{}_out_opt.rknn'.format(platform)
if __name__ == '__main__':add_perm = False # 如果设置成True,则将模型输入layout修改成NHWC# Create RKNN objectrknn = RKNN(verbose=True)# pre-process configprint('--> config model')rknn.config(batch_size=1, mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], reorder_channel='0 1 2', target_platform=[platform], force_builtin_perm=add_perm, output_optimize=1)print('done')# Load tensorflow modelprint('--> Loading model')ret = rknn.load_onnx(model=ONNX_MODEL,outputs=['output', '327', '328'])# 这里一定要根据onnx模型输出节点进行修改if ret != 0:print('Load onnx model failed!')exit(ret)print('done')# Build modelprint('--> Building model')ret = rknn.build(do_quantization=True, dataset='./dataset.txt')if ret != 0:print('Build rkmodel failed!')exit(ret)print('done')# rknn.export_rknn_precompile_model(RKNN_MODEL)rknn.export_rknn(RKNN_MODEL)rknn.release()
六、通过toolkit将onnx转换为rknn模型
1、将上述第二步导出的 yolov5s.onnx 模型以及 第五步创建的 onnx2rknn.py 放到服务器当前目录，即映射到docker容器内/test目录的目录；
2、进入docker容器内，执行
cp /test/yolov5s.onnx /examples/onnx/yolov5/
cp /test/onnx2rknn.py /examples/onnx/yolov5/
cd /examples/onnx/yolov5/
python onnx2rknn.py
cp yolov5s_rk3399pro_out_opt.rknn /test
执行完成后将生成相应的 yolov5s_rk3399pro_out_opt.rknn 模型
七、在rk3399pro开发板上运行rknn模型
1、下载 airockchip/RK3399Pro_npu，将其放到 rk3399pro 开发板上；
2、如有必要，根据说明更新NPU驱动，只需要替换相应的boot.img 即可；
3、进入 RK3399Pro_npu-main/rknn-api/examples/c_demos/rknn_yolov5_demo 目录，执行以下命令
cd RK3399Pro_npu-main/rknn-api/examples/c_demos/rknn_yolov5_demo
mkdir build
cd build
cmake -DCMAKE_C_COMPILER=/usr/bin/gcc -DCMAKE_CXX_COMPILER=/usr/bin/g++ ..
make -j4
执行成功后，将在build目录生成可执行程序 rknn_yolov5_demo
4、将上述第六步导出的yolov5s_rk3399pro_out_opt.rknn模型上传到rk3399pro的
RK3399Pro_npu-main/rknn-api/examples/c_demos/rknn_yolov5_demo/model 目录
5、在RK3399Pro_npu-main/rknn-api/examples/c_demos/rknn_yolov5_demo创建以下运行脚本 run_demo.sh
export LD_LIBRARY_PATH=./lib/
./build/rknn_yolov5_demo model/yolov5s_rk3399pro_out_opt.rknn model/bus.bmp
6、执行脚本
./run_demo.sh
正常将输出预测信息
可视化结果如图所示
