目标检测 YOLOv5 - 最新版本v6.2模型在瑞芯微 Rockchip设备上运行的方案
目标检测 YOLOv5 - 最新版本v6.2模型在瑞芯微 Rockchip设备上运行的方案
flyfish
模型的导出
本例是以640x640的输入来说明
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install
可以切换到v6.2然后执行导出onnx
python export.py --weights yolov5s.pt --include torchscript onnx
或者直接下载v6.2的源码
需要更改导出部分
这个为了说明如何编写代码,导出了4部分。在使用时要么输出3部分,要么输出1部分,只选其一。
方案1 模型的输出如下
[1,25200,85]
方案2 模型的输出如下
(1, 3, 80, 80, 85)
(1, 3, 40, 40, 85)
(1, 3, 20, 20, 85)
在写代码时根据输出的name编写
ret = rknn.load_onnx(model=ONNX_MODEL, outputs=['339', '377', '415'])
if ret != 0:
print('Load model failed!')
exit(ret)
print('done')
方案2完整的由onnx转rknn代码下载地址
https://github.com/shaoshengsong/rockchip_rknn_yolov5
在yolov5_v6.2文件夹中
方案3
需要在Ubuntu下进行预编译和量化
为了Rockchip的设备使用量化,这里把box和score分开了。
box范围是0-640
score的范围是0-1
box 由[x_center, y_center, width, height] 变成了[left, top, right, bottom]
常见的边框(bounding box )坐标表示方法
最后模型导出的样子如下,输出两部分