admin 管理员组

文章数量: 1184232

目标检测 YOLOv5

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

flyfish

源码地址


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版本的源码进行训练


模型导出

当一个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 CPUFusing layers... 
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients, 16.4 GFLOPsTorchScript: 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 installedONNX: 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:       

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

文件所在路径 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,后面跟着分数。

其他方案

本文标签: 目标检测 YOLOv5