YOLO v5 TensorRT C++推演快速教程

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

一些较为流行的模型推演方案包括Google的TensorFlow Serving，META (FaceBook) 的 TorchServe，Microsoft的ONNX Runtme，百度的Paddle Lite，华为的CANN，腾讯的TNN，阿里的MNN等。这些推演方案面向不同的部署场景与运算设备平台，力求最小化模型推演成本。

然而，这些现成的方案并不一定能满足全部应用场景。如果你的模型推演需要运行在搭载了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

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

上述依赖项是TensorRT开发的基本条件。除此之外，我假设你已经配置好了自己喜欢的集成开发环境。如果不确定的话，可以考虑VS Code，并记得安装C++插件。

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

在使用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侧存储（内存），然后执行后处理逻辑。

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

接下来，你将通过YOLO v5这个实例来掌握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

然后遵循官方教程，导出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&...

GitHubultralytics

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

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

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

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

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

TensorRT的CMake工程

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

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

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 中如下：

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 中如下：

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 引用和链接正确的头文件与库文件了。

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

YOLO v5的TensorRT推演代码

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

yolov5-trt/yolov5.cpp at main · yinguobing/yolov5-trt

YOLO v5 inference with TensorRT (C++). Contribute to yinguobing/yolov5-trt development by creating an account on GitHub.

GitHubyinguobing

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

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

辅助模块

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

// 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;

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

// 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检测结果的一个辅助结构体。

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

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

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

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

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

理解代码主体

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

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

我们逐个来看。

资源初始化

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

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

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

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

// 从磁盘读取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。

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

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

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 记录下与推演有关的内存区域。

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

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

// 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 。它用来在异步操作中同步结果。

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

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

输入图像读取与预处理

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

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

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

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;
    }
}

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

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

// 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;
}

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

执行模型推演

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

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

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

至此，模型推演完成。

数据后处理

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

// 分配输出内存空间: "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);
}

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

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

// 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;
    }
}

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

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

// 同步结果
cudaStreamSynchronize(stream);

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

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

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

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

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

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

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

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

    ...
    // 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);
    ...

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

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

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

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

    ...
    p += step;
}

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

// 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);

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

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

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

// 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.

GitHubyinguobing

祝推演顺利！