DETR tensorRT 的 C++ 部署

山水无移-ZhangQian

已于 2024-01-10 11:20:21 修改

阅读量765

点赞数 18

文章标签： c++ transformer 目标检测

于 2024-01-08 16:50:37 首次发布

本文链接：https://blog.csdn.net/zhangqian_1/article/details/135461014

版权

DETR tensorRT 的 C++ 部署

本篇说说DETR tensorRT 的 C++ 部署。

【完整代码、模型、测试图片】

1 导出 onnx 模型（建议先看）

方法1：导出DETR onnx并修改模型输出Gather层，解决tesorrt 推理输出结果全为0问题，参考【DETR TensorRT python接口部署】。

方法2：导出DETR模型更有效的onnx新方法，推荐使用【DETR tensor去除推理过程无用辅助头+fp16部署再次加速+解决转tensorrt 输出全为0问题的新方法】。

2 编译

编译前修改 CMakeLists.txt 对应的TensorRT版本
在这里插入图片描述

cd DETR_tensorRT_Cplusplus
mkdir build
cd build
cmake ..
make

3 运行

运行时如果.trt模型存在则直接加载，若不存会自动先将onnx转换成 trt 模型，并保存在给定的路径，然后运行推理

# 运行时如果.trt模型存在则直接加载，若不存会自动先将onnx转换成 trt 模型，并保存在给定的路径，然后运行推理。
cd build
./detr_trt

4 测试效果

onnx 测试效果
在这里插入图片描述
tensorRT 测试效果

tensorRT 时耗
使用的显卡 Tesla V100、cuda_11.0

5替换模型说明

1）导出的onnx模型建议simplify后，修改Gather层后再转trt模型。
2）注意修改后处理相关 postprocess.hpp 中相关的类别参数。

修改相关的路径

    std::string OnnxFile = "/zhangqian/workspaces1/TensorRT/detr_trt_Cplusplus/models/detr_r50_person_sim_change.onnx";
    std::string SaveTrtFilePath = "/zhangqian/workspaces1/TensorRT/detr_trt_Cplusplus/models/detr_r50_person_sim_change.trt";
    cv::Mat SrcImage = cv::imread("/zhangqian/workspaces1/TensorRT/detr_trt_Cplusplus/images/test.jpg");

    int img_width = SrcImage.cols;
    int img_height = SrcImage.rows;

    CNN DETR(OnnxFile, SaveTrtFilePath, 1, 3, 640, 640, 3);  // 1, 3, 640, 640, 3 前四个为模型输入的NCWH, 3为模型输出叶子节点的个数+1，（本示例中的onnx模型输出有2个叶子节点，再+1=3）
    DETR.ModelInit();
    DETR.Inference(SrcImage);

    for (int i = 0; i < DETR.DetectiontRects_.size(); i += 6)
    {
        int classId = int(DETR.DetectiontRects_[i + 0]);
        float conf = DETR.DetectiontRects_[i + 1];
        int xmin = int(DETR.DetectiontRects_[i + 2] * float(img_width) + 0.5);
        int ymin = int(DETR.DetectiontRects_[i + 3] * float(img_height) + 0.5);
        int xmax = int(DETR.DetectiontRects_[i + 4] * float(img_width) + 0.5);
        int ymax = int(DETR.DetectiontRects_[i + 5] * float(img_height) + 0.5);

        char text1[256];
        sprintf(text1, "%d:%.2f", classId, conf);
        rectangle(SrcImage, cv::Point(xmin, ymin), cv::Point(xmax, ymax), cv::Scalar(255, 0, 0), 2);
        putText(SrcImage, text1, cv::Point(xmin, ymin + 15), cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 0, 255), 2);
    }

    imwrite("/zhangqian/workspaces1/TensorRT/detr_trt_Cplusplus/images/result.jpg", SrcImage);

    printf("== obj: %d \n", int(float(DETR.DetectiontRects_.size()) / 6.0));