> ## Content Index
> Fetch the complete content index at: https://yinguobing.com/blog/llms.txt
> Use this file to discover other available public pages before exploring further.

# YOLO v5 TensorRT C++推演快速教程
- URL: https://yinguobing.com/blog/a-quick-tutorial-for-yolov5-tensorrt-cpp-inference/
- Published: 2022-08-13T13:24:29.000Z
- Updated: 2026-09-08T02:07:32.000Z
- Description: 经过数天乃至数周的调试，你的神经网络模型终于在速度与精度上获得了令人满意的结果，并正式进入模型部署方案开发阶段。如何与现有的产品耦合，并在目标设备上最大化模型推演性能，成为一项重要课题。
- Author: 尹国冰
- Tags: 深度学习

模型推演通常是指模型接收给定输入数据、完成内部逻辑运算并输出计算结果的过程。尤其是在模型部署阶段，不再需要反向传播计算梯度，推演成为模型的唯一使命。正因如此，推演阶段大概率不会也不需要再使用TensorFlow、PyTorch这样庞杂的框架，而视情况使用专门为模型推演设计的方案。

一些较为流行的模型推演方案包括Google的[TensorFlow Serving](https://github.com/tensorflow/serving?ref=yinguobing.com)，META (FaceBook) 的 [TorchServe](https://pytorch.org/serve/?ref=yinguobing.com)，Microsoft的[ONNX Runtme](https://onnxruntime.ai/?ref=yinguobing.com)，百度的[Paddle Lite](https://www.paddlepaddle.org.cn/?ref=yinguobing.com)，华为的[CANN](https://www.hiascend.com/software/cann?ref=yinguobing.com)，腾讯的[TNN](https://www.baidu.com/link?url=fQlpxb7cWx145Fav5DPJfVYTBP9ftqXSelnF6VS-V9yI6Abdt0hhT%5F6SrFO8Qbxd&wd=&eqid=f9394ccb00090a950000000662e62cc5&ref=yinguobing.com)，阿里的[MNN](https://www.baidu.com/link?url=GOoE%5FIZuKeaq8hc99ELzIpDDVYqaUWpa17qUlq5EUtxcvwsMX8y6GOJPDMjJxVsB&wd=&eqid=a29cf2f8000a8b190000000662e62d56&ref=yinguobing.com)等。这些推演方案面向不同的部署场景与运算设备平台，力求最小化模型推演成本。

然而，这些现成的方案并不一定能满足全部应用场景。如果你的模型推演需要运行在搭载了NVIDIA公司GPU的设备上，并且需要与已有的产品做整合，那么，更加底层的TensorRT或许是你的最佳选择。

TensorRT是由NVIDIA公司开发的一款用于在GPU上执行深度神经网络运算的代码库，专注于加速深度神经网络的运算，支持PyThon与C++。虽然在模型训练时Python被大量使用，但在产品基本定型进入最终开发阶段时，从最大化性能的角度考虑，C++是更加理想的选择。这篇文章以物体检测框架YOLO v5为例，介绍了如何使用TensorRT的C++ API实现一个完整的物体检测功能。

教程分为三部分：

1. 搭建开发环境，安装必要的依赖项。
2. 了解TensorRT开发涉及到的一些基本概念。
3. 动手操作，实战YOLO v5示例代码。

作为这篇教程的读者，我假设你已经具备基本的深度神经网络知识，并且有使用Python推演神经网络模型的代码经验。这篇快速教程将重点放在TensorRT的C++ API上，你将会看到使用C++ API完成模型推演与使用Python的一些显著不同之处。

准备好了吗，让我们开始吧！

（全文约5500字）

## 准备工作 

### 工欲善其事，必先利其器

⛳

本章节，你将会验证目标机器是否已经安装了开发所需的必要依赖项目。对于不满足的依赖项，文中给出了指导性的安装教程供你参考。

本教程内容在Ubuntu Server 20.04下测试通过。使用的GPU型号为NVIDIA RTX 3090。部分关键的依赖项目与版本号如下：

- NVIDIA驱动：470.82.01
- CUDA：11.4.3
- cuDNN：8.2.1
- TensorRT：8.2.5.1
- GCC：9.4.0
- Cmake：3.16.3

如果你的库版本低于上述值，最好将其升级到该版本或者更高版本。

#### ‌如何确定NVIDIA驱动版本

按下 `Ctrl + Alt + T` 打开终端，然后输入 `nvidia-smi` 即可显示驱动版本。该命令同时可以显示GPU的型号。

如果你需要安装新驱动，可参阅：

[https://yinguobing.com/install-nvidia-driver-runfile-ubuntu-server-20-04/](https://yinguobing.com/blog/install-nvidia-driver-runfile-ubuntu-server-20-04/)

#### 如何确定CUDA版本

在终端中输入 `nvcc --version` 命令，即可显示CUDA版本号。

如果你需要安装CUDA，可参阅：

[https://yinguobing.com/install-cuda11-with-runfile/](https://yinguobing.com/blog/install-cuda11-with-runfile/)

#### 如何确定cuDNN版本

在终端中输入：

`cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2`

即可显示cuDNN版本信息。

如果需要安装cuDNN，可参考：

[https://yinguobing.com/install-cudnn8-with-tar-file/](https://yinguobing.com/blog/install-cudnn8-with-tar-file/)

#### 如何确定TensorRT版本

在终端中输入 `find /usr -name NvInferVersion.h` 查找头文件路径，

然后解析头文件内容：

`cat /usr/local/TensorRT-8.2.5.1/include/NvInferVersion.h | grep NV_TENSORRT -A 2`

即可显示TensorRT版本信息。

如果需要安装TensorRT，可参考：

[https://yinguobing.com/install-tensorrt-with-tar-file/](https://yinguobing.com/blog/install-tensorrt-with-tar-file/)

#### 如何确定GCC版本

在终端中输入 `gcc --version` 即可显示版本信息。

如果需要安装 GCC，可以在终端中输入：

`sudo apt install gcc`

#### 如何确定CMake版本

在终端中输入`cmake --version` 即可显示版本信息。

如果需要安装 cmake，可以在终端中输入：

`sudo apt install cmake`

上述依赖项是TensorRT开发的基本条件。除此之外，我假设你已经配置好了自己喜欢的集成开发环境。如果不确定的话，可以考虑[VS Code](https://www.google.com/url?sa=t&rct=j&q=&esrc=s&source=web&cd=&cad=rja&uact=8&ved=2ahUKEwjBl62Z56L5AhU0tlYBHQB5DjIQFnoECCQQAQ&url=https%3A%2F%2Fcode.visualstudio.com%2F&usg=AOvVaw15O90sm1ios8AUpw56hCml)，并记得安装C++插件。

恭喜你，可以进入下一章节了。

## TensorRT基本概念

### TensorRT Basic Concepts

⛳

本章节你将了解TensorRT中的以下几个关键概念：Runtime、CudaEngine、ExecutionContext、与Binding。

在使用ONNX Runtime的Python API执行模型推演大致分四步：

1. 初始化：使用onnx格式存储的模型文件初始化 `InferenceSession` 。
2. 预处理：将输入数据做预处理。
3. 执行推演：调用 `run` 方法执行模型推演，获得模型输出。
4. 后处理：执行后处理逻辑。

‌TensorRT的推演与其类似，只是要更加复杂一些：

1. 模型准备：TensorRT使用私有模型格式，其它格式的模型文件如ONNX需要先完成转换才可以被使用。该转换过程可以通过调用TensorRT的C++或者Python API来在线实现，也可以通过 `trtexec` 这个工具离线完成。
2. 初始化：包括初始化运行时 `Runtime` ，使用 `Runtime` 加载模型文件形成 `CudaEngine` ，再基于 `CudaEngine` 创建 `ExecutionContext` ；这三者组合起来可以大致等同为 ONNX Runtime中的 `InferenceSession` 。此外还需要按照模型的输入与输出规格分别申请CPU一侧的内存与GPU一侧的显存，指向这些存储区域的指针集合称为 `Binding` 。
3. 将输入数据做预处理，然后将其从CPU侧存储（内存）拷贝到GPU侧存储（显存）。
4. 调用 `ExecutionContext` 的`enqueueV2` 方法异步执行模型推演，然后调用 `cudaStreamSynchronize` 方法同步结果，获得模型输出。
5. 将输出结果从GPU侧存储（显存）拷贝到CPU侧存储（内存），然后执行后处理逻辑。

💡

TensorRT支持同步与异步推演，这里以异步推演作为说明实例。

与ONNX Runtime Python API的推演流程相比，TensorRT C++ API推演明显有着更加底层、更加细粒度的数据与流程控制。这些正是使用C++作为开发语言与Python相比所带来的主要区别。

接下来，你将通过YOLO v5这个实例来掌握TensorRT的实际使用过程。本章节所涉及到的概念都会出现在代码中，希望到时候你看到它们时不会感觉太陌生。

## 实战YOLO V5推演

### Time to get your hands dirty!

⛳

本章节你将借助YOLO v5模型推演的完整代码掌握TensorRT C++ 推演API，以及将ONNX格式神经网络模型转换为TensorRT专用格式的具体方法。

YOLO v5是一款在物体检测领域颇为流行的神经网络模型。在这里我们不会深入探究其原理与设计，而将直接使用GitHub上的一款开源实现所生成的预训练模型。在此基础上，本章节展示了以下内容。

1. 使用TensorRT的 `trtexec` 工具将ONNX模型文件转换为TensorRT的Engine格式。
2. 使用CMake构建TensorRT工程。
3. 使用OpenCV实现输入图像读取与存储。
4. YOLO v5模型预处理与后处理。
5. 使用TensorRT实现YOLO v5模型推演。

让我们开始吧！

### TensorRT模型转换

首先下载YOLO v5预训练好的模型权重文件，其地址为：

[https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt](https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt?ref=yinguobing.com)

然后遵循官方教程，导出ONNX格式模型：

[TFLite, ONNX, CoreML, TensorRT Export · Issue #251 · ultralytics/yolov5📚 This guide explains how to export a trained YOLOv5 🚀 model from PyTorch to ONNX and TorchScript formats. UPDATED 7 July 2022\. Before You Start Clone repo and install requirements.txt in a Python&...![](https://github.com/fluidicon.png)GitHubultralytics![](https://opengraph.githubassets.com/c8b731ba2ebd757995b8286cd47302c4dea17bd9a1189c58b73fe85d53846342/ultralytics/yolov5/issues/251)](https://github.com/ultralytics/yolov5/issues/251?ref=yinguobing.com)

上边这一步完成后，你应当获得一个ONNX格式模型文件 `yolov5s.onnx` 。

接下来，使用如下命令将其转换为TensoRT格式模型文件。

```bash
trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s.engine
```

其中 `yolov5s.onnx` 为输入， `yolov5s.engine` 为输出。该指令将以默认选项生成TensorRT适用的模型文件。需要注意的是，TensoRT模型的转换操作隐含了模型优化过程，所以转换后的模型与转换时的硬件设备GPU、TensorRT等软件版本存在关联性。不匹配的硬件型号与软件版本号有可能给出错误的运算结果。因此在后续部署时切记要确认部署环境的软硬件兼容性。

💡

本文为快速教程故不对转换参数做细致的介绍，如果你感兴趣的话可以参考TensorRT的官方文档。

如果一切顺利，你现在应当获得了TensorRT模型文件 `yolov5s.engine` ，该文件将用于在内存中构建 `CudaEngine` 。在介绍这部分代码之前，我们还需要完成一项准备工作——TensorRT CMake工程

### TensorRT的CMake工程

TensorRT的安装仅仅是拷贝头文件与库文件。因此要让CMake正确生成构建用文件的前提是找到正确的文件路径。具体包括：CUDA头文件与库文件、TensorRT的头文件 `NvInfer.h` 以及库文件 `libnvinfer.so` 。

对于CUDA，可以使用 `find_package` 函数来查找。 在 `CMakeFileLists.txt` 中如下： 

```cmake
find_package(CUDA REQUIRED)
message(STATUS "Found CUDA headers: ${CUDA_INCLUDE_DIRS}")
message(STATUS "Found CUDA libraries: ${CUDA_LIBRARIES}")
```

对于TensorRT，暂不支持 `find_package` 函数，因此需要借助 `find_path` 函数来查找头文件。在 `CMakeFileLists.txt` 中如下：

```cmake
set(TENSORRT_ROOT /usr/local/TensorRT-8.2.5.1)
find_path(TENSORRT_INCLUDE_DIR NvInfer.h
        HINTS ${TENSORRT_ROOT} PATH_SUFFIXES include/)
message(STATUS "Found TensorRT headers: ${TENSORRT_INCLUDE_DIR}")
```

以及 `find_library` 函数来查找库文件。在 `CMakeFileLists.txt` 中如下：

```cmake
find_library(TENSORRT_LIBRARY_INFER nvinfer
        HINTS ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES lib lib64 lib/x64)
message(STATUS "Found TensorRT-nvinfer libs: ${TENSORRT_LIBRARY_INFER}")
```

其中 `TENSORRT_ROOT` 为TensorRT的实际安装目录。之后，便可以使用 `target_include_directories` 与 `target_link_libraries` 引用和链接正确的头文件与库文件了。

💡

受限于篇幅，这里不再列出完整的CMakeFileLists文件内容。你可以在文后给出的GitHub链接里找到完整的工程文件。

至此，全部准备工作完成，是时候展示代码了！

### YOLO v5的TensorRT推演代码

对于那些心急的人，你可以在这里找到完整的推演代码：

[yolov5-trt/yolov5.cpp at main · yinguobing/yolov5-trtYOLO v5 inference with TensorRT (C++). Contribute to yinguobing/yolov5-trt development by creating an account on GitHub.![](https://github.com/fluidicon.png)GitHubyinguobing![](https://opengraph.githubassets.com/31fae213c70249f36b34e69dbee4d3da16db07db9d661454a850561072706a6a/yinguobing/yolov5-trt)](https://github.com/yinguobing/yolov5-trt/blob/main/yolov5.cpp?ref=yinguobing.com)

YOLO v5 TensoRT C++推演代码

全部代码行数为249行——比起Python来说多了不少。不过不要担心，理解这些代码，需要的只是一点点耐心。

全部代码大致可以分为两部分：模型推演直接相关的 `run` 函数，以及其它辅助模块。从代码量上来看 `run` 函数共计199行，是绝对的主体部分。因此在顺序上我们先来了解下代码量较少的辅助模块，最后再集中精力理解 `run` 函数的内容。

### 辅助模块

辅助模块的第一部分是日志模块。这是一个必选模块，在构建TensorRT Runtime时需要作为参数传入。日志模块必须继承自 `nvinfer1::ILogger` 类。为了简单起见，这里引用GitHub用户[@cyrusbehr](https://github.com/cyrusbehr?ref=yinguobing.com)的实现如下：

```cpp
// A logger is REQUIRED for TensorRT
// Ref: https://github.com/cyrusbehr/tensorrt-cpp-api
using Severity = nvinfer1::ILogger::Severity;
class Logger : public nvinfer1::ILogger {
    void log(Severity severity, const char* msg) noexcept override;
};

void Logger::log(Severity severity, const char* msg) noexcept
{
    if (severity <= Severity::kWARNING) {
        std::cout << msg << std::endl;
    }
}

Logger logger;
```

TensorRT推演logger实现

💡

如果你需要一个更加完整的日志实现，可以参考NVIDIA的官方示例代码。

第二个模块是用于计算模型的输入与输出数据所占据的内存大小的函数，返回的数值单位为字节 (byte)。后续在申请内存时会用到。

```cpp
// A helper function to calculate memory usage
size_t get_memory_size(const nvinfer1::Dims& dims, const int32_t elem_size)
{
    return std::accumulate(dims.d, dims.d + dims.nbDims, 1, std::multiplies<int64_t>()) * elem_size;
}
```

用于计算内存区域大小的辅助函数

第三个模块为用于记录YOLO v5检测结果的一个辅助结构体。

```cpp
// A struct to store the detection results
struct Result {
    float score;
    cv::Rect box;
    int class_id;
};
```

用于存储YOLO检测结果的结构体

其中 `score` 为检测到物体的置信度， `box` 为物体的边界框， `class_id` 为物体的类别ID。

第四个模块为可执行文件的入口函数 `main` 。

```cpp
int main(int argc, char const* argv[])
{
    if (argc != 3) {
        std::cout << "Run like this:\n    " << argv[0] << "yolov5s.engine input.jpg" << std::endl;
        return -1;
    }

    run(argv[1], argv[2]);

    return 0;
}
```

入口 main 函数

`main` 函数的两个输入参数分别为TensorRT模型文件 `yolov5s.engine` ，以及待检测的输入图像 `input.jpg` 。这两个参数在运行时根据实际状况修改。

了解了这些辅助模块及其用途之后，我们可以集中精力来处理代码的主体部分—— `run` 函数了。

### 理解代码主体

`run` 函数的逻辑分为4部分：

1. 资源初始化，包括运行时、引擎、运行环境与显存等。
2. 输入图像读取与预处理。
3. 引擎执行推演。
4. 输出数据后处理。

我们逐个来看。

### 资源初始化

**推演运行时** `nvinfer1::IRuntime` 是一切资源的源头，是第一个需要被初始化的对象。

```cpp
std::unique_ptr<nvinfer1::IRuntime> runtime { nvinfer1::createInferRuntime(logger) };
```

初始化推演运行时

注意运行时的初始化需要将辅助模块中定义好的 `logger` 作为传入参数。同时这段代码使用 `unique_ptr` 来跟踪初始化后的运行时对象。

**引擎** `nvinfer1::ICudaEngine` 存储了模型的权重参数，它的初始化分2步：首先将模型文件读入内存，然后使用运行时的 `deserializeCudaEngine` 方法初始化引擎对象。

```cpp
// 从磁盘读取engine文件到内存
std::ifstream engine_file(engine_file_path, std::ios::binary);
if (engine_file.fail()) {
    std::cout << "Failed to read model file." << std::endl;
}
engine_file.seekg(0, std::ifstream::end);
auto fsize = engine_file.tellg();
engine_file.seekg(0, std::ifstream::beg);
std::vector<char> engineData(fsize);
engine_file.read(engineData.data(), fsize);

// 初始化引擎
std::unique_ptr<nvinfer1::ICudaEngine> mEngine { runtime->deserializeCudaEngine(engineData.data(), fsize) };
if (mEngine.get() == nullptr) {
    std::cout << "Failed to deserialize CUDA engine." << std::endl;
}

```

初始化引擎

引擎的本质是专为TensorRT优化之后的模型。然而模型在推演时，其内部参数参与运算之后所获得的数值也需要统一管理——这正是**推演环境**要做的事情。一旦引擎构建完成，便可以使用它的 `createExecutionContext` 方法来构建推演环境`nvinfer1::IExecutionContext`。

```cpp
std::unique_ptr<nvinfer1::IExecutionContext> context { mEngine->createExecutionContext() };
if (context.get() == nullptr) {
    std::cout << "Failed to create CUDA context." << std::endl;
}
```

创建推演环境

TensorRT的模型推演在GPU上进行，你需要在显存上为用于计算的数据申请空间。下面这段代码展示了如何为输入数据申请显存空间。

```cpp
int32_t channels = 3, height = 640, width = 640;
nvinfer1::Dims4 input_dims { 1, channels, height, width };
size_t input_mem_size = get_memory_size(input_dims, sizeof(float));
void* cuda_mem_input { nullptr };
if (cudaMalloc(&cuda_mem_input, input_mem_size) != cudaSuccess) {
    std::cout << "ERROR: input cuda memory allocation failed, size = " << input_mem_size << " bytes" << std::endl;
}
```

为输入申请显存空间

上边这段代码根据YOLO v5s的输入尺寸 3\*640\*640 以及辅助模块中的内存大小计算函数得到显存中应当开辟的空间大小，然后使用 `cudaMalloc` 申请显存。其中 `cuda_mem_input` 为指向该显存区域的指针。为了方便后续推演，这里使用一个指针数组 `bindings` 记录下与推演有关的内存区域。

```cpp
// Keep track of 1 input and 4 output device memory buffers
void* bindings[5];
bindings[0] = cuda_mem_input;
```

记录输入显存区域

依据同样的道理，以下代码为模型输出申请显存空间。

```cpp
// Allocate CUDA memory for output bindings
std::vector<std::string> output_node_names { "339", "392", "445", "output" };
std::vector<size_t> output_mem_sizes;
for (size_t i = 0; i < output_node_names.size(); i++) {
    int32_t output_idx = mEngine->getBindingIndex(output_node_names[i].c_str());
    if (output_idx == -1) {
        std::cout << "ERROR: failed to get output by name: " << output_node_names[i] << std::endl;
            break;
    }
    auto output_dims = context->getBindingDimensions(output_idx);
    auto output_size = get_memory_size(output_dims, sizeof(float));
    output_mem_sizes.push_back(output_size);
    void* cuda_mem_output { nullptr };
    if (cudaMalloc(&cuda_mem_output, output_size) != cudaSuccess) {
        std::cout << "ERROR: output cuda memory allocation failed, size = " << output_size << " bytes" << std::endl;
        break;
    } else {
        bindings[1 + i] = cuda_mem_output;
    }
}

```

申请输出显存空间的代码

这段代码更长一些，是因为模型的输出包含多个对象，必须为每个书对象分配独立的显存区域，为此使用了循环操作。在循环体中，首先通过 `getBindingDimensions` 获得输出每个对象的维度，然后计算其占据的显存区域大小，并调用 `cudaMalloc` 函数分配显存，最后将指向该存储区域的指针存入之前创建好的 `bindings` 。

显存空间初始化完成后，还需要初始化最后一个对象 `cudaStream` 。它用来在异步操作中同步结果。

```cpp
// Create a CUDA stream
cudaStream_t stream { nullptr };
if (cudaStreamCreate(&stream) != cudaSuccess) {
    std::cout << "ERROR: cuda stream creation failed." << std::endl;
}
```

创建 CudaStream 对象

终于，初始化操作全部完成。Python API下两三行可以搞定的事情，在这里代码行数翻了十倍不止。好在这些初始化之后的资源在推演过程中可以复用，而且考虑TensorRT所带来的性能增益，这些付出是值得的。

### 输入图像读取与预处理

OpenCV作为传统图像处理库，提供了大量图像处理方法。这里我们将使用它来将图像从文件读入到内存。注意这里的图像通道顺序为 `BGR` 。

```cpp
// Read in an image with OpenCV
cv::Mat img_bgr = cv::imread(image_file_path);
```

OpenCV 读取图像

输入图像在馈入神经网络之前，首先会按照一定规则做预处理操作。典型的预处理操作包括：缩放到指定尺寸、减去均值、归一化、除方差。另外如果模型对输入通道的顺序有要求，预处理也会调整通道顺序。YOLO v5的输入大小为 640x640像素，通道顺序为 `RGB` ， 预处理也非常简单：像素值除以255即可。下面这段代码展示了对应的预处理过程。

```cpp
cv::resize(img_bgr, img_bgr, cv::Size(width, height));
float input_buffer[height * width * channels] { 0 };
for (int c = 0; c < channels; c++) {
    for (int j = 0, HW = height * width; j < HW; ++j) {
        input_buffer[c * HW + j] = static_cast<float>(img_bgr.data[j * channels + 2 - c]) / 255.0f;
    }
}
```

YOLO v5预处理代码

其中 `img_bgr.data` 为指向像素存储区域的指针，而括号内下表运算中的 `+2` 实现了BGR数据到RGB的转换。

💡

从方便理解的角度出发，这里使用了两层嵌套循环实现全部像素的处理。在实际项目中，你应当寻求效率更高的做法。

数据预处理的操作是在HOST侧完成，在推演之前还需要将预处理后的数据拷贝到GPU一侧。

```cpp
// Memory copy: CPU-MEM to GPU-MEM
if (cudaMemcpyAsync(cuda_mem_input, input_buffer, input_mem_size, cudaMemcpyHostToDevice, stream) != cudaSuccess) {
    std::cout << "ERROR: CUDA memory copy of input failed, size = " << input_mem_size << " bytes" << std::endl;
}
```

将输入数据从HOST一侧拷贝到GPU一侧

至此，输入数据的准备工作全部完成，是时候执行模型推演了。

### 执行模型推演

虽然模型推演是最耗费算力的操作，但是，这部分的实际代码可能比你想的要简单。

```cpp
// 异步执行推演
bool status = context->enqueueV2(bindings, stream, nullptr);
if (!status) {
    std::cout << "ERROR: TensorRT inference failed." << std::endl;
}
```

这段代码中 `enqueueV2` 是真正触发模型推演的操作。参数中的 `bindings` 是我们之前记录的指向输入与输出显存的指针数组。 `stream` 是我们在初始化资源章节创建的 `cudaStream` 对象，用于同步数据。

💡

`enqueueV2` 是异步操作。如果你希望使用同步操作，请参见TensorRT官方文档。

至此，模型推演完成。

### 数据后处理

模型推演的下一步为输出数据后处理。首先按照输出数据的规格申请输出内存空间。

```cpp
// 分配输出内存空间: "339", "392", "445", "output"
std::vector<float*> output_buffers;
for (size_t i = 0; i < output_mem_sizes.size(); i++) {
    float* buf = new float[output_mem_sizes[i] / sizeof(float)];
    output_buffers.push_back(buf);
}
```

为输出申请HOST一侧内存空间

这段代码为四个输出结果申请了四份内存空间，并将指向这些内存区域的指针存储在 `output_buffers` 指针数组中。

然后，将模型推演获得的结果从GPU一侧显存拷贝到CPU一侧内存。

```cpp
// Memory copy: GPU-MEM to CPU-MEM
for (size_t i = 0; i < output_mem_sizes.size(); i++) {
    auto mem_to_host_result = cudaMemcpyAsync(output_buffers[i], bindings[1 + i], output_mem_sizes[i], cudaMemcpyDeviceToHost, stream);
    if (mem_to_host_result != cudaSuccess) {
        std::cout << "ERROR: CUDA memory copy of output " << i << " failed, size = " << output_mem_sizes[i] << " bytes" << std::endl;
    }
}
```

拷贝GPU数据到HOST一侧

这段代码中使用了 `cudaMemcpyAsync` 函数。其中参数 `output_buffers` 为刚才申请好的内存空间， `bindings` 为GPU一侧的现存空间， `output_mem_sizes` 为输出结果所占据的内存大小， `cudaMemcpyDeviceToHost` 指定了拷贝操作的方向，由于这是一个异步操作，所以要使用 `stream` 作为同步数据的依据。

到这一步，本教程已经调用了3个异步操作函数，因此这里需要使用同步操作来等待执行真正完成。

```cpp
// 同步结果
cudaStreamSynchronize(stream);
```

一旦这行代码执行结束，模型推演的结果才真正存储在了内存 `output_buffers` 中。接下来便是YOLO v5的后处理时刻。

💡

这部分代码仅对关键部分做说明，完整的代码请参阅GitHub连接。

首先定义三个向量分别存储检测获得的物体边界框、分值与类别ID。

```cpp
std::vector<float> scores;
std::vector<cv::Rect> boxes;
std::vector<int> class_ids;
```

存储结果的向量

YOLO v5模型有4个输出，但是我们需要的只是其中的 `output` 。 `output` 本身是一个巨大的矩阵，其行数为25200，代表模型给出的物体检测结果。为此，可以定义一个指针 `p` 指向 `output` 的第一个元素，然后通过一个循环来轮训结果。

```cpp
float* p = output_buffers[3];
int step = 85, proposal_count = 25200;
for (size_t i = 0; i < proposal_count; i++) {
    ...
```

后处理循环

每次循环获取结果包含85个元素，为物体的具体属性。其中第5个元素代表了物体检测的分值，可以用来过滤低置信度结果：

```cpp
    ...
    float obj_score = p[4];
    if (obj_score < score_threshold) {
        p += step;
        continue;
    }
    ...
```

获取检测分值

第6到85为物体类别ID的概率分布，可以按照最大值获取到分类结果：

```cpp
    ...
    // The type of the object?
    int c_id = -1;
    float c_score = 0;
    for (size_t j = 5; j < step; j++) {
        if (p[j] > c_score) {
            c_score = p[j];
            c_id = j - 5;
        }
    }
    scores.push_back(c_score * obj_score);
    ...
```

获取类别ID

注意YOLO v5检测到的类别分值 `c_score` 还需要乘以物体分值 `obj_score` 来获得最终的置信度结果。

第0-3为物体边界框：中心x，中心y，宽度与高度：

```cpp
    boxes.push_back(cv::Rect((p[0] - p[2] / 2) / scale, (p[1] - p[3] / 2) / scale,p[2] / scale, p[3] / scale));
```

获取物体的边界框

处理完一个物体之后，指针运算指向下一行结果：

```cpp
    ...
    p += step;
}
```

别忘记指针运算指向下一个待处理元素

如此重复操作直到全部输出被处理完成，获得初步过滤与转换后的结果：物体的位置、分值与类别。不过这还没完，接下来很重要的一步为非极大抑制NMS，用来融合高度叠加的重复结果。

```cpp
// Perform non maximum suppression to eliminate redundant overlapping boxes with lower confidences
std::vector<int> indices;
cv::dnn::NMSBoxes(boxes, scores, score_threshold, nms_threshold, indices);
```

NMS

这段代码使用了OpenCV提供的NMS方法 `cv::dnn::NMSBoxes` 。其中 `boxes` 为存储了 `Rect` 类的向量， `scores` 为存储了分值的向量， `score_threshold` 为检测分值阈值， `nms_threshold` 为NMS操作的阈值， `indices` 为一个空向量，用来存储NMS的结果。最后按照 `indices` 给出的物体序号提取最终结果。

```cpp
// Collect the detection result
std::vector<Result> results;
for (auto i : indices) {
    results.push_back(Result { scores[i], boxes[i], class_ids[i] });
}
```

按照NMS输出提取结果

最后，可以借助OpenCV将检测结果绘制在图像上，并保存到磁盘。

```cpp
// Draw all the results in the image
for (auto& r : results) {
    cv::rectangle(img_bgr, r.box, cv::Scalar(0, 255, 255), 2);
}

 // Save it
cv::imwrite("result.jpg", img_bgr);
```

绘制并保存结果

---

### 总结

以上为教程的全部内容。TensorRT C++的推演代码行数明显高于Python，不过大部分的代码都与准备工作相关，并且大量的涉及到了内存显存相关的操作。一旦熟悉了这些操作的背景，它们看起来也是很自然的过程。实际工程中，初始化工作可以封装起来，对外暴露一个友好的接口，让代码看起来更加简洁。

最后再次附上工程地址：

[GitHub - yinguobing/yolov5-trt: YOLO v5 inference with TensorRT (C++)YOLO v5 inference with TensorRT (C++). Contribute to yinguobing/yolov5-trt development by creating an account on GitHub.![](https://github.com/fluidicon.png)GitHubyinguobing![](https://opengraph.githubassets.com/31fae213c70249f36b34e69dbee4d3da16db07db9d661454a850561072706a6a/yinguobing/yolov5-trt)](https://github.com/yinguobing/yolov5-trt?ref=yinguobing.com)

祝推演顺利！