目标检测 YOLOv5 - v6.2版本模型在瑞芯微 Rockchip设备从训练到C++部署实践

目标检测 YOLOv5 - v6.2版本模型在瑞芯微 Rockchip设备从训练到C++部署实践

flyfish

源码地址

https://github.com/shaoshengsong/rockchip_rknn_yolov5

Rockchip 支持 YOLOv5 v6.2 从训练到C++部署的全链条开发，包括

C++部署代码
pytorch模型转onnx代码
onnx模型转rknn代码

文章目录

YOLOv5模型版本

模型导出环境
rknn-toolkit 1.7.1

部署环境
设备测试环境：RV1126
交叉编译环境：Ubuntu18.04

训练

从yolov5官网下载v6.2版本的源码进行训练

https://github.com/ultralytics/yolov5

模型导出

当一个yolov5_6.2模型训练完成后

1 pt模型 转 onnx

进入rockchip_rknn_yolov5\yolov5_6.2_export目录
执行

python export.py --weights yolov5s_v6.2.pt --img 640 --batch 1 --include onnx torchscript

yolov5s_v6.2.pt 是训练的模型，执行该命令后生成onnx模型
结果

export: data=data/coco128.yaml, weights=['yolov5s_v6.2.pt'], imgsz=[640], batch_size=1, device=cpu, half=False, inplace=False, train=False, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, include=['onnx','torchscript']
YOLOv5 🚀 2022-8-17 Python-3.6.9 torch-1.10.1+cu102 CPU

Fusing layers... 
YOLOv5s summary:213 layers,7225885 parameters,0 gradients,16.4 GFLOPs

TorchScript: starting export with torch 1.10.1+cu102...
TorchScript: export success, saved as yolov5s_v6.2.torchscript (27.8 MB)
Python 3.7.0 required by YOLOv5, but Python 3.6.9is currently installed

ONNX: starting export with onnx 1.11.0...
ONNX: export success, saved as yolov5s_v6.2.onnx (27.6 MB)

Export complete (3.78s)
Results saved to /home/a/deeplearning/original/yolov5-6.2
Detect:          python detect.py --weights yolov5s_v6.2.onnx 
Validate:        python val.py --weights yolov5s_v6.2.onnx 
PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5','custom','yolov5s_v6.2.onnx')
Visualize:       https://netron.app

2 onnx模型 转 rknn

进入rockchip_rknn_yolov5\yolov5_6.2_export目录
执行

python yolov562_to_rknn_3_4.py

yolov562_to_rknn_3_4的输出 ,此仓库使用该方式
最后的3_4表示有三个输出，每个输出4维

type: float32[1,255,80,80]type: float32[1,255,40,40]type: float32[1,255,20,20]

在执行前需要更为自己的配置
各个文件的路径，以下全提供在bin文件夹中

1 路径配置

ONNX_MODEL ='d:\\yolov5s_v6.2_output3_4.onnx'
RKNN_MODEL ='d:\\yolov5s_v6.2_output3_4.rknn'
IMG_PATH ='d:\\bus.jpg'
DATASET ='d:\\dataset.txt'

2 是否需要预编译

转换rknn模型代码是跨平台运行，需要预编译在Ubuntu上将 pre_compile=True (windows不支持)

 ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET, pre_compile=False)

3 连接到开发板

填写自己设备的型号和ID

ret = rknn.init_runtime('rv1126', device_id='bab4d7a824f04867')

执行结果

--> Config model
done
--> Loading model
done
--> Building model
done
--> Export RKNN model
done
--> Init runtime environment
done
--> Running model
class: person, score:0.8047849535942078
box coordinate left,top,right,down:[471.46940302848816,233.36096787452698,562.3111145496368,517.8007752895355]class: person, score:0.8034061789512634
box coordinate left,top,right,down:[210.58499628305435,238.00635361671448,284.57671636343,522.446161031723]class: person, score:0.7967345118522644
box coordinate left,top,right,down:[113.3626811504364,236.31811678409576,204.1182758808136,541.2366927862167]class: person, score:0.36264267563819885
box coordinate left,top,right,down:[78.74371695518494,336.6768593788147,123.03675484657288,517.1076865196228]class: bus , score:0.762096107006073
box coordinate left,top,right,down:[88.67809492349625,121.26825976371765,556.7039851546288,474.0234272480011]

生成rknn模型并看到测试结果

4 其他的onnx转rknn的方式

rockchip_rknn_yolov5\yolov5_6.2_export\other_export

yolov5_v6.2_python_output1_3表示模型输出是

（1,25200,85） 

yolov5_v6.2_python_output3_5表示输出是

(1,3,80,80,85)(1,3,40,40,85)(1,3,20,20,85)

中间转换的shape

(3,80,80,85)(3,40,40,85)(3,20,20,85)

最终使用的shape

(80,80,3,85)(40,40,3,85)(20,20,3,85)

C++ 部署

文件

YOLOv562Detector.cpp  
YOLOv562Detector.h  
test_image.cc

文件所在路径

rockchip_rknn_yolov5\C++\yolov5_62

使用方法

test_image "/userdata/yolov5s_v6.2_output3_4.rknn" 80 640

可执行文件名字 test_image
1 模型路径

/userdata/yolov5s_v6.2_output3_4.rknn

2 模型推理类别个数

80

3 模型输入大小

640

0表示人，5表示公共汽车，显示的是类别ID，后面跟着分数。

其他方案