Convert yolov5 PyTorch model to TensorRT

1. github open source code

The open source code location of yolov5 TensorRT inference is at https://github.com/linghu8812/tensorrt_inference/tree/master/yolov5 , the code of PyTorch to onnx can be found in the code from the original author's fork: https://github.com/linghu8812/ yolov5 , made a little modification to the model conversion.

2. Convert PyTorch model to ONNX model

First use the command git clone https://github.com/linghu8812/yolov5.gitclone yolov5 code, then copy the export_onnx.py file to the yolov5/modelsfolder, and generate the ONNX file through the following command. For yolov5s, yolov5m, yolov5l, yolov5x, these models can all be supported. --weightsYou can specify the path of the model file, --imgand --batchset the batch size for the input image size.

export PYTHONPATH="$PWD" && python3 export_onnx.py --weights ./weights/yolov5s.pt --img 640 --batch 1

With YOLOv4 models, the output of the concat also made, as shown in FIG.
yolov5
The onnxsim module is added to the code to simplify the structure of the subsequent tranpose and reshape nodes. The simplified code is as follows:

        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)

The model structure is simplified to:
Insert picture description here

3. Convert ONNX model to TensorRT model

3.1 Overview

The TensorRT model is the reasoning engine of TensorRT, and the code is implemented in C++. The relevant configuration is written in the config.yaml file, if engine_filethe path exists , it will be read engine_file, otherwise it will be onnx_filegenerated 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);
    }
}

The config.yaml file can set the batch size, the size of the image and the anchor of the model.

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 Compile

Compile the project with the following command to generateyolov5_trt

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

3.3 Operation

Run the project through the following command to get the inference result

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

4. Inference results

The inference result is shown in the figure below:
Insert picture description here

The above picture is the test result of yolov5x. In the output information, the average processing time of each picture of yolov5x is about 22.9ms, and the reasoning time of a single engine is about 12.1ms.