【深度学习模型高效部署】tensorrtx 深度解读：TensorRT 高性能推理实战案例

以下内容将对 tensorrtx 仓库做一个完整的介绍，包括其核心价值、主要功能、应用案例以及示例代码（附详细解释），帮助你快速上手并了解如何将深度学习模型高效部署在 NVIDIA GPU 上进行推理。

---

一、项目概述

tensorrtx 是由开发者 wang-xinyu 维护的一个针对 NVIDIA TensorRT 的示例项目集合。它提供了多种常用深度学习模型（主要涵盖 目标检测、图像分割、分类 等）的 TensorRT 推理实现示例。这些示例可以帮助开发者把在 PyTorch、TensorFlow 等深度学习框架下训练好的模型，快速转换并部署到 TensorRT 中，从而获得 低延迟、高吞吐量 的推理性能。

1.1 主要覆盖的模型及功能

• 目标检测：
  • YOLO 系列（YOLOv3、YOLOv4、YOLOv5、YOLOv6、YOLOv7、YOLOv8 等）
  • YOLOX
  • CenterNet
  • SSD
  • RetinaNet
• 分类：
  • ResNet、MobileNet 等
• 分割：
  • U-Net、SegFormer 等
• 其他常见网络:
  • DeTR、DBFace 等

1.2 核心价值

1. 简洁、直观的示例工程：
   每个模型都对应一个独立的文件夹和项目配置，涵盖 模型转换（通常是 .pth / .onnx → TensorRT engine）、推理代码、后处理 等完整流程。

2. 高效部署参考：
   tensorrtx 提供了最佳实践示例，许多优化技巧（如 INT8/FP16、Plugins、Cuda Stream 等）在示例中都有体现，能帮助你在实际项目中快速验证和应用。

3. 可持续更新：
   该项目随 TensorRT 和 各大主流网络版本 的迭代而不断更新，紧跟前沿模型（如 YOLOv8）的部署需求。

4. 适配多种平台：
   tensorrtx 一般在 Linux + NVIDIA GPU 环境下使用。也有部分开发者把它迁移到 Windows 或 Jetson 设备上，通过 CMake 配置可灵活切换。

---

二、使用流程概览

以 YOLOv5 为例来简述 tensorrtx 的大体使用流程（大多数模型类似）：

1. 准备训练好的模型权重

   • 例如在 PyTorch 中训练好的 best.pt 或者从官方提供的模型权重下载 yolov5s.pt。

2. 导出 ONNX

   • 使用官方或者 tensorrtx 提供的脚本，将 PyTorch 格式的 .pt 文件导出为 ONNX 格式，比如 yolov5s.onnx。

3. 生成 TensorRT engine

   • 在 tensorrtx 的 yolov5 目录下，有一个 yolov5.cpp 或者 yolov5_gen_trt.cpp 用于将 .onnx 转化为 TensorRT engine（比如 yolov5s.engine）。
   • 主要依赖 TensorRT C++ API，也可能需要一些 Plugin。

4. 推理

   • 利用生成的 yolov5s.engine，在 C++ 代码中通过 TensorRT API 构建 ICudaEngine 和 IExecutionContext，再将图像数据输入，引擎输出的结果经过后处理得到检测框/类别等。

5. 验证

   • 根据测试图片或视频，运行推理可视化结果，评估速度和精度。

---

三、示例应用案例

以下以 YOLOv5 为代表，介绍一个完整的示例。其他网络（YOLOv4、YOLOv7 等）在 tensorrtx 中的用法非常类似，差别主要在于 导出 ONNX 的脚本和 后处理 的代码不同。

3.1 文件结构

tensorrtx/yolov5 目录典型结构如下（可能随版本略有不同）：

tensorrtx
 └── yolov5
      ├── CMakeLists.txt           // 编译配置
      ├── yolov5.cpp              // 推理主程序(可生成引擎并推理)
      ├── yololayer.cu            // YOLO后处理插件(Anchor、NMS等)
      ├── ...
      └── README.md                // 使用说明

3.2 关键编译步骤

1. 安装 TensorRT（需与 GPU 驱动、CUDA、CUDNN 匹配）

2. 克隆 tensorrtx

   git clone https://github.com/wang-xinyu/tensorrtx.git
   cd tensorrtx/yolov5
   

3. 准备 ONNX

   • 比如将 yolov5s.pt 转换成 yolov5s.onnx。官方脚本示例：

     # 在 yolov5 官方仓库中，如 ultralytics/yolov5
     python export.py --weights yolov5s.pt --include onnx
     # 生成 yolov5s.onnx
     

4. 编译并运行

   mkdir build && cd build
   cmake ..
   make -j
   # 生成可执行文件 yolov5
   ./yolov5 -s yolov5s.onnx yolov5s.engine    # 先序列化生成engine
   ./yolov5 -d yolov5s.engine                # 进行推理测试
   

   • 选项 -s 表示 serialize（将 onnx 转为 engine 文件）
   • 选项 -d 表示 inference（使用 engine 文件进行推理测试）

---

四、核心代码分析

下面以一个 简化版 的 YOLOv5 推理代码（C++）为例，帮助你理解其核心逻辑。该示例演示如何在 C++ 中使用 TensorRT 构建引擎并执行推理。注意：这是一个简化示例，实际 tensorrtx 中包含完整的 Plugin、NMS、Anchor 等实现。

4.1 生成引擎示例

#include <NvInfer.h>
#include <iostream>
#include <fstream>
#include <memory>

// Helper: 将文件读取到内存buffer中
std::vector<char> readFile(const std::string& fileName) {
    std::ifstream file(fileName, std::ios::binary | std::ios::ate);
    if (!file.good()) {
        std::cerr << "Error: could not open file " << fileName << std::endl;
        return {};
    }
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<char> buffer(size);
    file.read(buffer.data(), size);
    file.close();
    return buffer;
}

// 生成 TensorRT engine
void createEngine(const std::string& onnxModelPath, const std::string& enginePath) {
    // 1. 创建构建器
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
    // 2. 创建网络定义
    const auto explicitBatch = 1U << static_cast<int>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
    // 3. 创建解析器
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    nvinfer1::IParser* parser = createParser(*network, gLogger);

    // 4. 解析 ONNX
    auto onnxData = readFile(onnxModelPath);
    if (onnxData.empty()) {
        std::cerr << "Failed to read ONNX file." << std::endl;
        return;
    }
    if (!parser->parse(onnxData.data(), onnxData.size())) {
        std::cerr << "Failed to parse ONNX." << std::endl;
        return;
    }

    // 5. 设置最大工作空间等配置
    config->setMaxWorkspaceSize(1ULL << 30); // 1GB
    // 可选：设置FP16/INT8等

    // 6. 构建引擎
    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    if (!engine) {
        std::cerr << "Failed to create engine." << std::endl;
        return;
    }

    // 7. 序列化引擎并写入文件
    nvinfer1::IHostMemory* serializedEngine = engine->serialize();
    std::ofstream p(enginePath, std::ios::binary);
    p.write(reinterpret_cast<const char*>(serializedEngine->data()), serializedEngine->size());
    p.close();

    // 8. 释放资源
    serializedEngine->destroy();
    engine->destroy();
    parser->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
}

int main(int argc, char** argv) {
    std::string onnxFile = "yolov5s.onnx";
    std::string engineFile = "yolov5s.engine";
    createEngine(onnxFile, engineFile);
    std::cout << "Engine created successfully!" << std::endl;
    return 0;
}

核心解释

1. 读取 ONNX 文件：将模型文件加载到内存中。
2. 使用 IBuilder、INetworkDefinition、IParser：这是 TensorRT 的标准流程，先创建网络结构，再将 ONNX 的计算图解析进网络。
3. 构建引擎：buildEngineWithConfig 会根据配置（batch size、精度模式等）将网络编译为一个高度优化的推理引擎。
4. 序列化：将内存中的引擎序列化为二进制数据，并写入到 .engine 文件。

---

4.2 推理示例

#include <NvInfer.h>
#include <cuda_runtime_api.h>
#include <opencv2/opencv.hpp>
#include <iostream>
#include <fstream>
#include <memory>

// 从文件读取引擎到内存并反序列化
nvinfer1::ICudaEngine* loadEngine(const std::string& engineFilePath) {
    std::ifstream file(engineFilePath, std::ios::binary);
    if (!file.good()) {
        std::cerr << "Error reading engine file: " << engineFilePath << std::endl;
        return nullptr;
    }
    file.seekg(0, file.end);
    size_t size = file.tellg();
    file.seekg(0, file.beg);
    std::vector<char> engineData(size);
    file.read(engineData.data(), size);
    file.close();

    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger);
    if (!runtime) {
        std::cerr << "Failed to create Infer Runtime." << std::endl;
        return nullptr;
    }
    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), size, nullptr);
    if (!engine) {
        std::cerr << "Failed to deserialize CUDA engine." << std::endl;
        return nullptr;
    }
    runtime->destroy();
    return engine;
}

int main(int argc, char** argv) {
    // 1. 加载 TensorRT engine
    std::string enginePath = "yolov5s.engine";
    nvinfer1::ICudaEngine* engine = loadEngine(enginePath);
    if (!engine) {
        return -1;
    }

    // 2. 创建执行上下文
    nvinfer1::IExecutionContext* context = engine->createExecutionContext();
    if (!context) {
        std::cerr << "Failed to create execution context." << std::endl;
        return -1;
    }

    // 假设网络输入大小为 (1,3,640,640)，请根据实际情况修改
    int inputIndex = engine->getBindingIndex("images"); // YOLOv5 的输入通常命名为 "images"
    int outputIndex = engine->getBindingIndex("output"); // YOLOv5 的输出名称可能不一样，需要根据实际模型查询

    // 3. 准备输入数据（以一张图片为例）
    cv::Mat img = cv::imread("bus.jpg"); // 测试图片
    cv::Mat resized;
    cv::resize(img, resized, cv::Size(640, 640)); 
    resized.convertTo(resized, CV_32FC3, 1/255.0f);

    // HWC -> CHW
    std::vector<float> inputData(3 * 640 * 640);
    int index = 0;
    for (int c = 0; c < 3; c++) {
        for (int h = 0; h < 640; h++) {
            for (int w = 0; w < 640; w++) {
                inputData[index++] = resized.at<cv::Vec3f>(h, w)[c];
            }
        }
    }

    // 4. 分配 GPU 内存
    void* buffers[2];
    cudaMalloc(&buffers[inputIndex], 3 * 640 * 640 * sizeof(float));
    // YOLO 输出大小不固定，简单起见这里先假设其为某固定值
    // 实际可通过engine->getBindingDimensions()获取
    int outputSize = 1000; // 这里只是示例
    cudaMalloc(&buffers[outputIndex], outputSize * sizeof(float));

    // 5. 将数据拷贝到 GPU
    cudaMemcpy(buffers[inputIndex], inputData.data(), 3 * 640 * 640 * sizeof(float), cudaMemcpyHostToDevice);

    // 6. 执行推理
    context->enqueueV2(buffers, 0, nullptr);

    // 7. 拷贝输出到 CPU
    std::vector<float> outputData(outputSize);
    cudaMemcpy(outputData.data(), buffers[outputIndex], outputSize * sizeof(float), cudaMemcpyDeviceToHost);

    // 8. 后处理（基于 YOLO 的输出做解码、NMS、可视化等）
    // 由于 YOLOv5 的输出格式比较复杂，这里省略细节。

    // 9. 释放资源
    cudaFree(buffers[inputIndex]);
    cudaFree(buffers[outputIndex]);
    context->destroy();
    engine->destroy();

    return 0;
}

关键步骤解读

1. 加载引擎并反序列化：用 IRuntime->deserializeCudaEngine() 将 .engine 文件载入内存中，得到 ICudaEngine 实例。
2. 创建执行上下文：IExecutionContext 用于真正执行推理。
3. 准备输入：
   • 图像预处理：大小、归一化、形状转换（NCHW）。
   • 拷贝到 GPU：TensorRT 推理时，数据需要放在 GPU 内存中。
4. 执行推理：
   • context->enqueueV2()（异步）或 context->executeV2()（同步）。
5. 后处理：
   • 对 YOLO 输出进行解析，得到检测框的坐标、类别、置信度，执行 NMS 筛选等。
6. 资源管理：
   • 及时销毁 context、engine，释放 CUDA 内存，避免内存泄漏。

---

五、优势与实际使用建议

1. 面向高性能推理场景：
   如果你的项目需要 实时或准实时 推断、批量处理，TensorRT 能够大大提升吞吐量并降低延迟。

2. 丰富的模型支持：
   tensorrtx 已经囊括了主流的目标检测、分割、分类网络，能快速用来做迁移学习部署或二次开发。

3. 插件化（Plugin）及自定义层：
   如果你有一些网络包含自定义算子，可以参考 tensorrtx 中的 Plugin 写法来扩展。

4. 跨平台能力：

   • 虽然大多数示例在 Ubuntu + NVIDIA GPU 测试，但也可以在 Windows 上用 CMake 进行编译，或在 Jetson 上进行部署（需要调整编译环境）。
   • tensorrtx 没有复杂的 Python 依赖，对于C++部署非常适合。

5. 持续更新：

   • tensorrtx 会随 YOLO 新版本或 TensorRT 新版本更新，及时提供可用的示例。

---

六、总结

tensorrtx 通过集合大量主流网络在 TensorRT 上的部署示例，为开发者提供了一个快速、直观、可复用的高性能推理参考项目。它的核心价值在于：

1. 降低学习门槛：把训练好的 ONNX 模型转换为 TensorRT engine 并进行高效推理，提供了标准化的示例；
2. 一站式优化：包含 FP16/INT8、Plugin、自定义算子等高级功能的示例，让你快速在项目中落地；
3. 紧跟前沿：支持 YOLOv8 等最新网络，满足不断演进的实际需求。

如果你正寻找将深度学习模型部署到 NVIDIA GPU 并获得高性能推理的解决方案，tensorrtx 无疑是一个非常值得参考和借鉴的开源项目。

---

引用链接

• tensorrtx GitHub 仓库: https://github.com/wang-xinyu/tensorrtx
• YOLOv5: https://github.com/ultralytics/yolov5

以上内容希望能帮助你对 tensorrtx 的整体架构、使用流程及示例代码有一个比较清晰的认识，能够更顺利地完成高性能推理部署。祝开发顺利！

【哈佛博后带小白玩转机器学习】 哔哩哔哩_bilibili

总课时超400+，时长75+小时