YOLO|Ubuntu18.04配置darknet环境实现YOLOv4目标检测(五)——darknet YOLOv4和YOLOv4-tiny模型转ONNX转TensorRT部署 TensorRT|计算机视觉|c++|深度学习

文章目录

1. github开源代码
2. darknet模型转ONNX模型
3. ONNX模型转TensorRT模型
- 3.1 概述
- 3.2 编译
- 3.3 运行
4. 推理结果

1. github开源代码开源代码位置在这里，darknet转ONNX模型代码基于python，TensorRT推理代码基于C++。
2. darknet模型转ONNX模型通过export_onnx.py文件可以将darknet模型转换为ONNX模型，目前可以支持YOLOv3，YOLOv3-SPP，YOLOv4等模型。

parser = argparse.ArgumentParser(description='Transform YOLO weights to ONNX.') parser.add_argument('--cfg_file', type=str, default='yolov4.cfg', help='yolo cfg file') parser.add_argument('--weights_file', type=str, default='yolov4.weights', help='yolo weights file') parser.add_argument('--output_file', type=str, default='yolov4.onnx', help='yolo onnx file') parser.add_argument('--neck', type=str, default='PAN', help='use which kind neck') args = parser.parse_args()

cfg_file是darknet的cfg文件的路径，weights_file是darknet的weights文件的路径，output_file是输出的ONNX文件的名称，neck是YOLO的模型结构，YOLOv3相关模型输入FPN，YOLOv4模型相关输入PAN。输入以下命令后，可以将darknet模型转为ONNX模型。

python3 export_onnx.py

【YOLO|Ubuntu18.04配置darknet环境实现YOLOv4目标检测(五)——darknet YOLOv4和YOLOv4-tiny模型转ONNX转TensorRT部署】导出yolov4-tiny模型的命令是：

python3 export_onnx.py --cfg_file yolov4-tiny.cfg --weights_file yolov4-tiny.weights --output_file yolov4-tiny.onnx --strides 32 16 --neck FPN

转换后的模型加入了transpose，将所有YOLO层的输出结果进行了concat，最终结构如下：

YOLO|Ubuntu18.04配置darknet环境实现YOLOv4目标检测(五)——darknet YOLOv4和YOLOv4-tiny模型转ONNX转TensorRT部署

文章图片

3. ONNX模型转TensorRT模型 3.1 概述 TensorRT模型即TensorRT的推理引擎，代码中通过C++实现。相关配置写在config.yaml文件中，如果存在engine_file的路径，则读取engine_file，否则从onnx_file生成engine_file。

void YOLOv5::LoadEngine() { // create and load engine std::fstream existEngine; existEngine.open(engine_file, std::ios::in); if (existEngine) { readTrtFile(engine_file, engine); assert(engine != nullptr); } else { onnxToTRTModel(onnx_file, engine_file, engine, BATCH_SIZE); assert(engine != nullptr); } }

config.yaml文件可以设置batch size，图像的size及模型的anchor等。

Yolov4: onnx_file:"../yolov4.onnx" engine_file:"../yolov4.trt" labels_file:"../coco.names" BATCH_SIZE:1 INPUT_CHANNEL: 3 IMAGE_WIDTH:608 IMAGE_HEIGHT:608 obj_threshold: 0.4 nms_threshold: 0.45 stride:[8, 16, 32] anchors:[[12, 16], [19, 36], [40, 28], [36, 75], [76, 55], [72, 146], [142, 110], [192, 243], [459, 401]]

3.2 编译通过以下命令对项目进行编译，生成Yolov4_trt