Yolov5 pytorch模型转tensorrt

yolov5 模型转tensorrt

一. github开源代码

yolov5 TensorRT推理的开源代码位置在https://github.com/linghu8812/onnx_trt_deploy/tree/master/yolov5,PyTorch转onnx的代码见从原作者fork过来的代码:https://github.com/linghu8812/yolov5,对模型转换做了一点修改。

二. pytorch模型转onnx模型

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

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

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

yolov5

三. ONNX模型转TENSORRT模型

3.1 概述

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

void YOLOv4::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);
    }
}

onfig.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. 推理结果

推理结果如下图所示:

result.jpg

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

本文链接:https://blog.csdn.net/linghu8812/article/details/109322729

已标记关键词 清除标记
©️2020 CSDN 皮肤主题: 大白 设计师:CSDN官方博客 返回首页