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

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

导出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，最终结构如下：

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

mkdir build && cd build
cmake ..
make -j

3.3 运行

通过以下命令运行项目，得到推理结果

./Yolov4_trt ../config.yaml ../samples

4. 推理结果

推理结果如下图所示：

可以看到输出结果与darknet模型的输出结果包括置信度完全一致，在输出的信息中，每张图片的平均处理时间约为16ms，单独engine的推理时间约为7ms。