yolov7 PyTorch模型转TensorRT

yolov7 PyTorch模型转TensorRT

1. github开源代码

yolov7 TensorRT推理的开源代码位置在https://github.com/linghu8812/tensorrt_inference/tree/master/project/yolov7，yolov7官方的开源代码位置为https://github.com/WongKinYiu/yolov7，对yolov7 epxort ONNX模型的代码作了少许修改，主要是(1) 在导出ONNX模型时，增加了onnxsim进行模型简化；(2) 在导出ONNX模型时进行了少许修改，可以去掉ONNX模型的ScatterND的op，使得TensorRT可以直接解析ONNX模型。

2. PyTorch模型转ONNX模型

首先通过命令git clone https://github.com/linghu8812/yolov7.gitclone yolov7的代码，通过以下命令生成ONNX文件。--weights 可以指定模型文件路径，--simplify进行ONNX模型简化，--grid使输出为一个tensor，--img-size为输入图片尺寸，--batch-size设置batch size的大小。

python export.py --weights ./weights/yolov7.pt --simplify --grid

对于1280输入的模型，可通过以下命令生成ONNX模型：

python export.py --weights ./weights/yolov7-w6.pt --simplify --grid --img-size 1280

3. ONNX模型转TensorRT模型

3.1 概述

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

void yolov7::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的处理已经写到了ONNX模型中，所以处理时只需知道每个检测层的stride和anchor的数量即可，不需要再配置具体的anchor。

yolov7:
    onnx_file:     "../yolov7.onnx"
    engine_file:   "../yolov7.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]
    num_anchors:   [3,  3,  3]

yolov7-w6模型的配置文件如下：

yolov7:
    onnx_file:     "../yolov7-w6.onnx"
    engine_file:   "../yolov7-w6.trt"
    labels_file:   "../coco.names"
    BATCH_SIZE:    1
    INPUT_CHANNEL: 3
    IMAGE_WIDTH:   1280
    IMAGE_HEIGHT:  1280
    obj_threshold: 0.4
    nms_threshold: 0.45
    strides:       [8, 16, 32, 64]
    num_anchors:   [3,  3,  3,  3]

3.2 编译

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

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

3.3 运行

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

./yolov7_trt ../config.yaml ../samples

4. 推理结果

推理结果如下图所示：