yolov5 PyTorch模型转TensorRT

yolov5 PyTorch模型转TensorRT

1. github开源代码

yolov5 TensorRT推理的开源代码位置在https://github.com/linghu8812/tensorrt_inference/tree/master/project/yolov5，PyTorch转onnx的代码见从原作者fork过来的代码：https://github.com/linghu8812/yolov5，对模型转换做了一点修改。
针对视频检测加跟踪的代码已开源，

• 开源地址在：https://github.com/linghu8812/tensorrt_tracker
• 视频运行结果：https://www.bilibili.com/video/BV1qg411K74p?t=0.0
• 博客介绍：https://blog.csdn.net/linghu8812/article/details/121178381

2. PyTorch模型转ONNX模型

首先通过命令git clone https://github.com/linghu8812/yolov5.gitclone yolov5的代码，通过以下命令生成ONNX文件。对于yolov5s，yolov5m，yolov5l，yolov5x这几个模型都可以支持。--weights 可以指定模型文件路径，--imgsz为输入图片尺寸，--batch-size设置batch size的大小。

python3 export.py ---weights weights/yolov5s.pt --batch-size 10 --imgsz 640 --include onnx --simplify

同YOLOv4模型一样，对输出结果也做了concat，如下图所示。

代码中加入了onnxsim模块，以简化后续的tranpose和reshape节点的结构，简化代码如下：

        onnx_model = onnx.load(f)  # load onnx model
        model_simp, check = simplify(onnx_model)
        assert check, "Simplified ONNX model could not be validated"
        onnx.save(model_simp, f)

模型结构简化为：

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等。

yolov5:
    onnx_file:     "../yolov5x.onnx"
    engine_file:   "../yolov5x.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
    stride:        [8, 16, 32]
    anchors:       [[10,13], [16,30], [33,23], [30,61], [62,45], [59,119], [116,90], [156,198], [373,326]]

3.2 编译

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

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

3.3 运行

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

./yolov5_trt ../config.yaml ../samples

4. 推理结果

推理结果如下图所示：

上图为yolov5x的测试结果，在输出的信息中，yolov5x每张图片的平均处理时间约为22.9ms，单独engine的推理时间约为12.1ms。