YOLOv6 PyTorch模型转TensorRT

1. github开源代码

YOLOv6 TensorRT推理的开源代码位置在https://github.com/linghu8812/tensorrt_inference/tree/master/YOLOv6，YOLOv6官方的开源代码位置为https://github.com/meituan/YOLOv6，页面上可以找到如何转换onnx模型和下载官方训练的COCO目标检测模型和官方已转好的onnx模型。

2. PyTorch模型转ONNX模型

首先通过命令git clone https://github.com/meituan/YOLOv6.gitclone YOLOv6的代码，通过以下命令生成ONNX文件。--weights 可以指定模型文件路径，--img-size为输入图片尺寸，--batch-size设置batch size的大小。

python deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1
1

也可以直接下载官方已转换好的onnx模型进行转换，下载地址为：https://github.com/meituan/YOLOv6/releases/tag/0.1.0

3. ONNX模型转TensorRT模型

3.1 概述

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

void YOLOv6::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);
    }
}
1
2
3
4
5
6
7
8
9
10
11
12

config.yaml文件可以设置batch size，图像的size，因为YOLOv6是anchor free目标检测模型，所以不需要设置anchor。

YOLOv6:
    onnx_file:     "../yolov6s.onnx"
    engine_file:   "../yolov6s.trt"
    labels_file:   "../coco.names"
    BATCH_SIZE:    1
    INPUT_CHANNEL: 3
    IMAGE_WIDTH:   640
    IMAGE_HEIGHT:  640
    obj_threshold: 0.4
    nms_threshold: 0.45
    strides:       [8, 16, 32]
1
2
3
4
5
6
7
8
9
10
11

3.2 编译

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

mkdir build && cd build
cmake ..
make -j
1
2
3

3.3 运行

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

./YOLOv6_trt ../config.yaml ../samples
1

4. 推理结果

推理结果如下图所示：

上图为YOLOv6s的测试结果，在输出的信息中，YOLOv6s每张图片的推理时间约为2.2ms左右，因为是anchor free的模型，所以与yolov5模型相比，YOLOv6的后处理时间可以减少一半。