TensorRT C++ API 教程

如何使用 TensorRT C++ API 实现高性能 GPU 机器学习推理。
支持具有单个/多个输入以及单个/多个输出的模型，并且可以进行批处理。

项目概述视频 . 代码深度解析视频

寻求维护者 🚀

本项目正在积极寻找维护者，以帮助指导其发展和改进。如果您对该项目充满热情并有意贡献，我非常期待与您交流！

欢迎通过 LinkedIn 联系我，讨论您如何参与其中。

TensorRT C++ 教程

我已阅读所有 NVIDIA TensorRT 文档，这样您就不必再费心了！

本项目演示如何使用 TensorRT C++ API 对图像数据进行高性能 GPU 推理。它涵盖了以下内容：

• 如何在 Ubuntu 20.04 / 22.04 上安装 TensorRT 10。
• 如何为您的 GPU 生成优化后的 TensorRT 引擎文件。
• 如何指定一个简单的优化配置文件。
• 如何运行 FP32、FP16 或 INT8 精度的推理。
• 如何从 GPU 内存读取/写入数据，并处理 GPU 图像。
• 如何使用 CUDA 流执行异步推理并在稍后同步。
• 如何处理具有静态和动态批大小的模型。
• 如何处理具有单个或多个输出张量的模型。
• 如何处理具有多个输入的模型。
• 包含一个 视频教程，我在其中逐行解释代码。
• 此代码可用作任何以固定尺寸图像/图像集作为输入的模型的基础，包括 Insightface ArcFace、YoloV8、SCRFD 人脸检测等。
  • 您只需实现相应的后处理代码即可。
• 待办事项：增加对输入形状动态变化的模型的支持。
• 待办事项：增加对 Windows 的支持。

入门指南

以下说明假定您正在使用 Ubuntu 20.04 或 22.04。 您需要提供自己的 ONNX 模型用于此示例代码，或者您可以下载示例模型（请参阅下方的“健康检查”部分）。

前提条件

• 已在 Ubuntu 20.04 和 22.04 上测试并正常工作（目前 不 支持 Windows）
• 安装 CUDA 11 或 12，安装说明 在此。
  • 推荐 >= 12.0
  • 必需 >= 11.0
• 安装 cuDNN，安装说明 在此。
  • 必需 >= 8
  • 必需 < 9（OpenCV GPU 尚不支持）
• sudo apt install build-essential
• sudo snap install cmake --classic
• sudo apt install libspdlog-dev libfmt-dev（用于日志记录）
• 安装带有 CUDA 支持的 OpenCV。要从源代码编译 OpenCV，请运行 ./scripts/ 中提供的 build_opencv.sh 脚本。
  • 如果您使用提供的脚本，并且将 cuDNN 安装到了非标准位置，则必须修改脚本中的 CUDNN_INCLUDE_DIR 和 CUDNN_LIBRARY 变量。
  • 推荐 >= 4.8
• 从 此处 下载 TensorRT 10。
  • 必需 >= 10.0
• 导航到 CMakeLists.txt 文件，将 TODO 替换为您 TensorRT 的安装路径。

构建库

• mkdir build
• cd build
• cmake ..
• make -j$(nproc)

运行可执行文件

• 导航到构建目录
• 运行可执行文件，并提供您的 ONNX 模型路径。
• 例如：./run_inference_benchmark --onnx_model ../models/yolov8n.onnx
  • 注意：有关如何获取 yolo8n 模型的说明，请参阅下方的“健康检查”部分。
• 首次针对特定模型和选项运行可执行文件时，将根据您的 ONNX 模型构建一个 TensorRT 引擎文件。此过程相对较慢，某些模型（如 YOLO 模型）可能需要 5 分钟以上。
• 或者，您也可以直接提供自己的 TensorRT 引擎文件：
  • 例如：./run_inference_benchmark --trt_model ../models/yolov8n.engine.NVIDIAGeForceRTX3080LaptopGPU.fp16.1.1
  • 注意：有关提供自定义 TensorRT 引擎文件时的注意事项，请参阅下方的 V5.0 更改日志。

健康检查

• 要进行健康检查，可从 这里 下载 YOLOv8n 模型。
• 接下来，使用以下脚本将其从 PyTorch 转换为 ONNX：
  • 您需要先运行 pip3 install ultralytics。

from ultralytics import YOLO
model = YOLO("./yolov8n.pt")
model.fuse()
model.info(verbose=False)  # 打印模型信息
model.export(format="onnx", opset=12) # 使用 opset 12 将模型导出为 ONNX

• 将生成的 ONNX 模型 yolov8n.onnx 放入 ./models/ 目录中。
• 使用该模型和位于 ./inputs/team.jpg 的图像进行推理，应产生以下特征向量：
  • 注意：由于 TensorRT 不是确定性的，因此特征向量不会完全相同，但会非常相似。

3.41113 16.5312 20.8828 29.8984 43.7266 54.9609 62.0625 65.8594 70.0312 72.9531 ...

INT8 推理

启用 INT8 精度可以在牺牲一定精度（由于动态范围减小）的情况下进一步加速推理。对于 INT8 精度，用户必须提供能够代表模型实际输入数据的校准数据。建议使用 1000 张以上的校准图像。要使用 YoloV8 检查模型启用 INT8 推理，需执行以下步骤：

• 在 main.cpp 中将 options.precision = Precision::FP16; 改为 options.precision = Precision::INT8;。
• 必须在 main.cpp 中修改 options.calibrationDataDirectoryPath = "";，以指定包含校准数据的路径。
  • 如果使用 YoloV8 模型，建议使用 COCO 验证数据集，可通过 wget http://images.cocodataset.org/zips/val2017.zip 下载。
• 确保 engine.cpp 中 Int8EntropyCalibrator2::getBatch 方法中的调整大小代码（见 TODO）适用于您的模型。
  • 如果使用 YoloV8 模型，预处理代码是正确的，无需更改。
• 重新编译并运行可执行文件。
• 校准缓存将被写入磁盘（.calibration 扩展名），以便在后续模型优化中重复使用。如果希望重新生成校准数据，则必须删除此缓存文件。
• 如果出现“函数 allocate 内存不足”的错误，则需要降低 Options.calibrationBatchSize 的值，以确保整个批次能够容纳在 GPU 内存中。

基准测试

基准测试在 RTX 3050 Ti 笔记本 GPU 和第 11 代 Intel(R) Core(TM) i9-11900H @ 2.50GHz 处理器上运行。

模型	精度	批量大小	平均推理时间
yolov8n	FP32	1	4.732 ms
yolov8n	FP16	1	2.493 ms
yolov8n	INT8	1	2.009 ms
yolov8x	FP32	1	76.63 ms
yolov8x	FP16	1	25.08 ms
yolov8x	INT8	1	11.62 ms

示例集成

想知道如何将此库集成到您的项目中吗？或者如何读取 YoloV8 模型的输出以提取有意义的信息？如果是这样，请查看我的两个最新项目：YOLOv8-TensorRT-CPP 和 YOLOv9-TensorRT-CPP，它们演示了如何使用 TensorRT C++ API 运行 YoloV8/9 推理（支持目标检测、语义分割和人体姿态估计）。这些项目在后端都使用了本项目！

项目结构

project-root/
├── include/
│   ├── engine/
│   │   ├── EngineRunInference.inl
│   │   ├── EngineUtilities.inl
│   │   └── EngineBuildLoadNetwork.inl
│   ├── util/...
│   ├── ...
├── src/
|   ├── ...
│   ├── engine.cpp
│   ├── engine.h
│   └── main.cpp
├── CMakeLists.txt
└── README.md

理解代码

• 大部分实现位于 include/engine 目录中。我在代码中添加了大量的注释，应该能帮助您轻松理解代码的工作原理。
• 推理代码位于 include/engine/EngineRunInference.inl。
• 构建和加载 TensorRT 引擎文件的代码位于 include/engine/EngineBuildLoadNetwork.inl。
• 您也可以观看我的深度解析视频，其中我逐行解释了所有代码。

调试方法

• 实现中使用了 spdlog 库进行日志记录。可以通过设置环境变量 LOG_LEVEL 为以下值之一来更改日志级别：trace、debug、info、warn、error、critical 或 off。
• 如果在从 ONNX 模型创建 TensorRT 引擎文件时遇到问题，可以尝试将环境变量 LOG_LEVEL 设置为 trace 并重新运行应用程序。这将为您提供有关构建过程具体失败位置的更多信息。

表达感谢

如果这个项目对您有所帮助，欢迎您给它点个赞。这将鼓励我保持项目的更新，并快速解决遇到的问题。如果您需要更具体的帮助，我也提供咨询服务。欢迎通过 LinkedIn 与我联系。

贡献者

Loic Tetrel 💻

thomaskleiven 💻

WiCyn 💻

更改日志

V6.0

• 现在实现要求 TensorRT >= 10.0。

V5.0

• Engine 类已修改为接受一个模板参数，用于指定模型的输出数据类型。现在支持 float、__half、int8_t、int32_t、bool 和 uint8_t 类型的输出。
• 增加了直接加载 TensorRT 引擎文件的支持，无需从 ONNX 模型编译。然而，强烈建议使用提供的 API 从 ONNX 模型构建引擎文件，而不是直接加载 TensorRT 模型。如果选择直接加载 TensorRT 模型文件，必须手动检查 Options 是否已正确设置（例如，如果您的模型是为 FP32 编译的，但您尝试运行 FP16 推理，则会失败，且可能不会给出详细的错误信息）。
• 添加了命令行解析器。

V4.1

• 增加了对固定批大小 > 1 的支持。

V4.0

• 增加了对 INT8 精度的支持。

V3.0

• 实现已更新为使用 TensorRT 8.6 API（例如 IExecutionContext::enqueueV3()）。
• 可执行文件已从 driver 重命名为 run_inference_benchmark，现在必须通过命令行参数传递 ONNX 模型路径。
• 移除了 Options.doesSupportDynamicBatchSize。实现现在会自动检测支持的批大小。
• 移除了 Options.maxWorkspaceSize。实现现在在模型构建过程中不限制 GPU 内存，允许使用尽可能多的内存池来存储中间层数据。

v2.2

• 将模型名称序列化为引擎文件的一部分。

V2.1

• 增加了对多输入模型的支持。现在支持单输入、多输入、单输出、多输出以及批处理的模型。

V2.0

• 需要安装 OpenCV CUDA。安装方法请参考 这里。
• 移除了 Options.optBatchSizes，替换为 Options.optBatchSize。
• 支持具有多个输出的模型（例如 SCRFD）。
• 增加了对不支持批处理推理的模型的支持（第一维输入固定）。
• 增加了更多的错误检查。
• 修复了用户在原始 V1.0 版本中经常遇到的一些常见问题。
• 移除了 GPU 设备名称中的空格。

贡献者 ✨

感谢以下各位优秀的贡献者（emoji key)：

本项目遵循 all-contributors 规范。欢迎任何形式的贡献！

TensorRT C++ API 快速上手指南

本指南基于 tensorrt-cpp-api 项目，帮助开发者在 Ubuntu 系统上快速构建并使用 TensorRT C++ API 进行高性能 GPU 推理。

环境准备

系统要求

• 操作系统: Ubuntu 20.04 或 22.04 (暂不支持 Windows)
• GPU: 支持 CUDA 的 NVIDIA 显卡

前置依赖

请确保安装以下组件（版本要求严格，请勿随意更改）：

1. CUDA Toolkit: 推荐 >= 12.0 (最低 11.0)
   • 下载地址：NVIDIA CUDA Downloads
2. cuDNN: 8.x 版本 (必须 < 9.0，因 OpenCV GPU 支持限制)
   • 下载地址：NVIDIA cuDNN
3. 基础编译工具:

   sudo apt install build-essential
   sudo snap install cmake --classic
   sudo apt install libspdlog-dev libfmt-dev
   

4. OpenCV (带 CUDA 支持):
   • 推荐版本 >= 4.8
   • 建议使用项目提供的脚本编译源码以启用 CUDA：

     # 进入项目 scripts 目录运行
     ./scripts/build_opencv.sh
     

   • 注意: 若 cuDNN 安装在非标准路径，需修改脚本中的 CUDNN_INCLUDE_DIR 和 CUDNN_LIBRARY 变量。
5. TensorRT:
   • 必须版本 >= 10.0
   • 下载地址：NVIDIA TensorRT 10
   • 下载后解压，并记录安装路径。

安装步骤

1. 配置项目

克隆项目代码后，打开根目录下的 CMakeLists.txt 文件，找到 TODO 占位符，将其替换为你本地的 TensorRT 安装路径。

2. 编译库与可执行文件

在项目根目录执行以下命令：

mkdir build
cd build
cmake ..
make -j$(nproc)

3. 准备测试模型 (Sanity Check)

若没有自己的 ONNX 模型，可使用 YOLOv8n 进行测试：

1. 安装 ultralytics:

   pip3 install ultralytics
   

2. 转换模型为 ONNX (Python 脚本):

   from ultralytics import YOLO
   model = YOLO("./yolov8n.pt")
   model.fuse()
   model.info(verbose=False)
   model.export(format="onnx", opset=12)
   

3. 将生成的 yolov8n.onnx 放入项目的 ./models/ 目录。
4. 确保 ./inputs/ 目录下有一张测试图片（如 team.jpg）。

基本使用

运行推理

进入 build 目录，执行以下命令加载 ONNX 模型并自动构建 TensorRT 引擎：

./run_inference_benchmark --onnx_model ../models/yolov8n.onnx

说明：首次运行时会自动从 ONNX 构建 .engine 文件，耗时约 5 分钟以上，后续运行将直接加载引擎文件。

指定精度 (可选)

默认使用 FP16。若需测试 INT8 加速：

1. 修改 src/main.cpp：
   • 将 options.precision = Precision::FP16; 改为 Precision::INT8;
   • 设置 options.calibrationDataDirectoryPath 为校准数据目录（如 COCO val2017 数据集）。
2. 重新编译并运行。

调试模式

若构建失败，可通过设置环境变量获取详细日志：

export LOG_LEVEL=trace
./run_inference_benchmark --onnx_model ../models/yolov8n.onnx