目标检测 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代码
文章目录
-
- 训练
- 模型导出
-
- 1 pt模型 转 onnx
- 2 onnx模型 转 rknn
-
- 1 路径配置
- 2 是否需要预编译
- 3 连接到开发板
- 4 其他的onnx转rknn的方式
- C++ 部署
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.9 is 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，后面跟着分数。
其他方案