yolov5使用TensorRT进行c++部署

---

前言

有了前面用c++进行opencv里dnn部署和onnxruntime部署的经验，使用TensorRT进行部署，我们只要了解tensorrt和cuda的一些相关api的使用即可方便的部署，整个部署流程都差不多。

1.安装tensorrt

官方网站下载和cuda，cudnn（可以高）对应的版本：
https://developer.nvidia.com/nvidia-tensorrt-download
进官网下载要先注册，然后注意下载对应的版本：

打上勾：

按照官网的安装步骤选择你下载的压缩包安装方式安装即可：
https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html

同样为了能正确使用tensorrt，一定要在相应的c++编译软件进行对应的头文件和依赖库文件配置。同时用到了cuda_runtime.h，也要检查下依赖库文件路径的配置

2.使用tensorrt进行c++部署

通过前面的单文件部署熟悉了整个流程后，直接进行相应的修改就行。同时，为了c++代码的习惯性和简洁性，进行分文件编写。

1.创建yolov5_trt.h头文件

头文件

// yolov5使用tensorrt进行部署的头文件

#ifndef YOLOV5_TRT_H
#define YOLOV5_TRT_H

#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <NvInfer.h>       // nvidia加载模型进行推理的插件
#include <NvOnnxParser.h>
#include <cuda_runtime.h>

一些自定义的结构

和上篇onnxruntime部署一样

// 自定义配置结构
struct Configuration
{
    
     
	float confThreshold; // Confidence threshold
	float nmsThreshold;  // Non-maximum suppression threshold
	float objThreshold;  //Object Confidence threshold
	std::string modelpath;
};

// 定义BoxInfo结构类型
typedef struct BoxInfo
{
    
    
	float x1;
	float y1;
	float x2;
	float y2;
	float score;
	int label;
} BoxInfo;

yolov5类模块

class YOLOv5
{
    
    
public:
	YOLOv5(Configuration config);
	~YOLOv5();  
	void UnInit();  //析构时调用，释放指针所指内存空间
	void detect(cv::Mat& frame);
private:
	float confThreshold;
	float nmsThreshold;
	float objThreshold;
	int inpWidth;
	int inpHeight;
	std::string classes[80] = {
    
    "person", "bicycle", "car", "motorbike", "aeroplane", "bus",
							"train", "truck", "boat", "traffic light", "fire hydrant",
							"stop sign", "parking meter", "bench", "bird", "cat", "dog",
							"horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe",
							"backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
							"skis", "snowboard", "sports ball", "kite", "baseball bat",
							"baseball glove", "skateboard", "surfboard", "tennis racket",
							"bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl",
							"banana", "apple", "sandwich", "orange", "broccoli", "carrot",
							"hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant",
							"bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse",
							"remote", "keyboard", "cell phone", "microwave", "oven", "toaster",
							"sink", "refrigerator", "book", "clock", "vase", "scissors",
							"teddy bear", "hair drier", "toothbrush"};

	const bool keep_ratio = true;
	void nms(std::vector<BoxInfo>& input_boxes);  
	cv::Mat resize_image(cv::Mat srcimg, int *newh, int *neww, int *top, int *left);

	void loadOnnx(const std::string strName);  // 第一次先加载onnx构建成engine
    void loadTrt(const std::string strName);  //  之后推理时都是将构建好的engine直接使用
	nvinfer1::ICudaEngine *m_CudaEngine;     // 在构建的网络上执行推理的引擎，具有功能上不安全的特征
	nvinfer1::IRuntime *m_CudaRuntime;		 // 允许对序列化的功能不安全的引擎进行反序列化
    nvinfer1::IExecutionContext *m_CudaContext;  // 使用具有功能不安全特性的引擎执行推理的上下文。
	cudaStream_t m_CudaStream;              // CUDA流的类型为cudaStream_t 
	int m_iInputIndex;						// 输入索引
    int m_iOutputIndex; 					
	int m_iClassNums;						// 类别数
    int m_iBoxNums;							// 输出预测的数量
	cv::Size m_InputSize;					// 输入尺寸
	void* m_ArrayDevMemory[2]{
    
    0};		   // 指向GPU内存的指针数组
    void* m_ArrayHostMemory[2]{
    
    0}; 
	int m_ArraySize[2]{
    
    0};
	std::vector<cv::Mat> m_InputWrappers{
    
    }; 
};

#endif

2.yolov5_trt.cpp文件的实现

头文件

// yolov5进行tensorrt部署的源文件
#include <fstream>
#include <iostream>
#include <sys/stat.h>
#include <glog/logging.h>
#include "./logging.h"
#include "yolov5_trt.h"

// 命名空间
using namespace cv;
using namespace nvinfer1;

推理部署的步骤

（1）加载用于tensorrt推理的模型文件

在有参构造的时候，加载onnx文件，将onnx文件构建为tensorrt的IR中间表示。因为从onnx开始加载，这个过程比较慢，所以首次构建成IR中间表示的时候，可以用一个文件保存构建的格式，这样再次加载的时候会直接从构建好的文件加载。（注意序列化引擎不能跨平台或TensorRT版本移植）。

这里我根据文档和代码总结，举例说明生成中间表示的过程并最终生成推理所需的上下文对象：
文档：https://docs.nvidia.com/deeplearning/tensorrt/archives/tensorrt-723/developer-guide/index.html

• 创建ILogger类型的全局对象。它是TensorRTAPI的各种方法的必需参数。下面是一个示例，演示如何创建记录器：
  这是https://github.com/onnx/onnx-tensorrt/blob/main/onnx_trt_backend.cpp的一个日志子类，可以自己写一些关于tensorrt的日志

// Logger for TRT info/warning/errors, https://github.com/onnx/onnx-tensorrt/blob/main/onnx_trt_backend.cpp
class TRT_Logger : public nvinfer1::ILogger
{
    
    
    nvinfer1::ILogger::Severity _verbosity;
    std::ostream* _ostream;

public:
    TRT_Logger(Severity verbosity = Severity::kWARNING, std::ostream& ostream = std::cout)
        : _verbosity(verbosity)
        , _ostream(&ostream)
    {
    
    
    }
    void log(Severity severity, const char* msg) noexcept override
    {
    
    
        if (severity <= _verbosity)
        {
    
    
            time_t rawtime = std::time(0);
            char buf[256];
            strftime(&buf[0], 256, "%Y-%m-%d %H:%M:%S", std::gmtime(&rawtime));
            const char* sevstr = (severity == Severity::kINTERNAL_ERROR ? "    BUG" : severity == Severity::kERROR
                        ? "  ERROR"
                        : severity == Severity::kWARNING ? "WARNING" : severity == Severity::kINFO ? "   INFO"
                                                                                                   : "UNKNOWN");
            (*_ostream) << "[" << buf << " " << sevstr << "] " << msg << std::endl;
        }
    }
};

• 使用C++API创建网络定义

	//创建生成器和网络
	IBuilder* builder = createInferBuilder(gLogger);    
	builder->setMaxBatchSize(1);   // batchsize
	const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);  // 显式批处理
	INetworkDefinition* network = builder->createNetworkV2(explicitBatch);   // 创建网络模型

• 使用解析器导入onnx模型并解析模型

    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);              // 使用nvonnxparser 定义一个可用的onnx解析器
    parser->parseFromFile(strModelName.c_str(), static_cast<int>(ILogger::Severity::kWARNING));   // 解析onnx

• 构建Engine

	// 使用builder对象构建engine
    IBuilderConfig* config = builder->createBuilderConfig();   // 
	// 特别重要的属性是最大工作空间大小
    config->setMaxWorkspaceSize(1ULL << 30);                   // 分配内存空间
    m_CudaEngine = builder->buildEngineWithConfig(*network, *config);    // 来创建一个 ICudaEngine 类型的对象，在构建引擎时，TensorRT会复制权重

• 序列化模型

	std::string strTrtName = strModelName;
    size_t sep_pos = strTrtName.find_last_of(".");
    strTrtName = strTrtName.substr(0, sep_pos) + ".trt";  // 
    IHostMemory *gieModelStream = m_CudaEngine->serialize();    // 将引擎序列化
    std::string serialize_str;     // 
    std::ofstream serialize_output_stream;
    serialize_str.resize(gieModelStream->size()); 
	// memcpy内存拷贝函数 ，从源内存地址的起始位置开始拷贝若干个字节到目标内存地址中
    memcpy((void*)serialize_str.data(),gieModelStream->data(),gieModelStream->size()); 
    serialize_output_stream.open(strTrtName.c_str());  
    serialize_output_stream<<serialize_str;     // 将引擎序列化数据转储到文件中
    serialize_output_stream.close();

• 最终生成推理所需的上下文对象

m_CudaContext = m_CudaEngine->createExecutionContext();    //执行上下文用于执行推理

整个（1）过程可以用一个函数实现：

// 加载onnx文件
void YOLOv5::loadOnnx(const std::string strModelName)
{
    
    
    Logger gLogger;   // 日志
    //根据tensorrt pipeline 构建网络
    IBuilder* builder = createInferBuilder(gLogger);    // 
    builder->setMaxBatchSize(1);   // batchsize
    const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);  // 显式批处理
    INetworkDefinition* network = builder->createNetworkV2(explicitBatch);                      // 定义模型
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);              // 使用nvonnxparser 定义一个可用的onnx解析器
    parser->parseFromFile(strModelName.c_str(), static_cast<int>(ILogger::Severity::kWARNING));   // 解析onnx
	// 使用builder对象构建engine
    IBuilderConfig* config = builder->createBuilderConfig();   // 
	// 特别重要的属性是最大工作空间大小
    config->setMaxWorkspaceSize(1ULL << 30);                   // 分配内存空间
    m_CudaEngine = builder->buildEngineWithConfig(*network, *config);    // 来创建一个 ICudaEngine 类型的对象，在构建引擎时，TensorRT会复制权重

    std::string strTrtName = strModelName;
    size_t sep_pos = strTrtName.find_last_of(".");
    strTrtName = strTrtName.substr(0, sep_pos) + ".trt";  // 
    IHostMemory *gieModelStream = m_CudaEngine->serialize();    // 将引擎序列化
    std::string serialize_str;     // 
    std::ofstream serialize_output_stream;
    serialize_str.resize(gieModelStream->size()); 
	// memcpy内存拷贝函数 ，从源内存地址的起始位置开始拷贝若干个字节到目标内存地址中
    memcpy((void*)serialize_str.data(),gieModelStream->data(),gieModelStream->size()); 
    serialize_output_stream.open(strTrtName.c_str());  
    serialize_output_stream<<serialize_str;     // 将引擎序列化数据转储到文件中
    serialize_output_stream.close();   
    m_CudaContext = m_CudaEngine->createExecutionContext();    //执行上下文用于执行推理
	// 使用一次，记得销毁parser，network, builder, and config 
    parser->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
}

再次加载的时候会直接从构建好的文件加载：

扫描二维码关注公众号，回复： 14858418 查看本文章

• 创建一个Runtime对象来反序列化:

    m_CudaEngine = m_CudaRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(), nullptr); // runtime对象反序列化

再次加载可以用另一个函数实现：

void YOLOv5::loadTrt(const std::string strName)
{
    
    
    Logger gLogger;
	// 序列化引擎被保留并保存到文件中
    m_CudaRuntime = createInferRuntime(gLogger);    
    std::ifstream fin(strName);
    std::string cached_engine = "";
    while (fin.peek() != EOF)
    {
    
     
        std::stringstream buffer;   
        buffer << fin.rdbuf();
        cached_engine.append(buffer.str());
    }
    fin.close();
    m_CudaEngine = m_CudaRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(), nullptr); // runtime对象反序列化
    m_CudaContext = m_CudaEngine->createExecutionContext();  //可以查询引擎获取有关网络的输入和输出的张量信息--维度/数据格式/数据类型
    m_CudaRuntime->destroy();
}

（2）获取输入输出节点的相关信息并保存

• 创建一些空间来存储中间激活值。由于引擎具有网络定义和训练参数，因此需要额外的空间。这些保存在上面通过中间表示生成的推理的上下文中

m_CudaContext = m_CudaEngine->createExecutionContext();  //可以查询引擎获取有关网络的输入和输出的张量信息--维度/数据格式/数据类型

• 使用输入和输出blob名称获取相应的输入和输出索引和相关信息：

	// 利用加载的模型获取输入输出信息
	// 使用输入和输出blob名来获取输入和输出索引
    m_iInputIndex = m_CudaEngine->getBindingIndex("images");     // 输入索引
    m_iOutputIndex = m_CudaEngine->getBindingIndex("output");   // 输出  
	Dims dims_i = m_CudaEngine->getBindingDimensions(m_iInputIndex);  // 输入，
	int size1 = dims_i.d[0] * dims_i.d[1] * dims_i.d[2] * dims_i.d[3];   // 展平
	m_InputSize = cv::Size(dims_i.d[3], dims_i.d[2]);   // 输入尺寸(W,H)
    Dims dims_o = m_CudaEngine->getBindingDimensions(m_iOutputIndex);  // 输出，维度[0,1,2,3]NHWC
	int size2 = dims_o.d[0] * dims_o.d[1] * dims_o.d[2];   // 所有大小
    m_iClassNums = dims_o.d[2] - 5;    // [,,classes+5]
    m_iBoxNums = dims_o.d[1];    // [b,num_pre_boxes,classes+5]

• 使用这些索引，设置指向GPU上输入和输出缓冲区的缓冲区阵列：

	// 分配内存大小
    cudaMalloc(&m_ArrayDevMemory[m_iInputIndex], size1 * sizeof(float));
    m_ArrayHostMemory[m_iInputIndex] = malloc(size1 * sizeof(float));
    m_ArraySize[m_iInputIndex] = size1 *sizeof(float);
    cudaMalloc(&m_ArrayDevMemory[m_iOutputIndex], size2 * sizeof(float));
    m_ArrayHostMemory[m_iOutputIndex] = malloc( size2 * sizeof(float));
    m_ArraySize[m_iOutputIndex] = size2 *sizeof(float);``

整个（1）和（2）可以利用构造函数在初始化阶段完成：

// 初始化
YOLOv5::YOLOv5(Configuration config)
{
    
    
	confThreshold = config.confThreshold;
	nmsThreshold = config.nmsThreshold;
	objThreshold = config.objThreshold;
	inpHeight = 640;
	inpWidth = 640;

	std::string model_path = config.modelpath;  // 模型权重路径
	// 加载模型
	std::string strTrtName = config.modelpath;      // 加载模型权重
	size_t sep_pos = model_path.find_last_of(".");
    strTrtName = model_path.substr(0, sep_pos) + ".engine"; // ".trt"
	if(ifFileExists(strTrtName.c_str()))
    {
    
            
        loadTrt(strTrtName);
    }
    else
    {
    
    
        loadOnnx(config.modelpath);
    }
	// 利用加载的模型获取输入输出信息
	// 使用输入和输出blob名来获取输入和输出索引
    m_iInputIndex = m_CudaEngine->getBindingIndex("images");     // 输入索引
    m_iOutputIndex = m_CudaEngine->getBindingIndex("output");   // 输出  
	Dims dims_i = m_CudaEngine->getBindingDimensions(m_iInputIndex);  // 输入，
	int size = dims_i.d[0] * dims_i.d[1] * dims_i.d[2] * dims_i.d[3];   // 展平
	m_InputSize = cv::Size(dims_i.d[3], dims_i.d[2]);   // 输入尺寸(W,H)
    Dims dims_o = m_CudaEngine->getBindingDimensions(m_iOutputIndex);  // 输出，维度[0,1,2,3]NHWC
	size = dims_o.d[0] * dims_o.d[1] * dims_o.d[2];   // 所有大小
    m_iClassNums = dims_o.d[2] - 5;    // [,,classes+5]
    m_iBoxNums = dims_o.d[1];    // [b,num_pre_boxes,classes+5]

	// 分配内存大小
    cudaMalloc(&m_ArrayDevMemory[m_iInputIndex], size * sizeof(float));
    m_ArrayHostMemory[m_iInputIndex] = malloc(size * sizeof(float));
    m_ArraySize[m_iInputIndex] = size *sizeof(float);
    cudaMalloc(&m_ArrayDevMemory[m_iOutputIndex], size * sizeof(float));
    m_ArrayHostMemory[m_iOutputIndex] = malloc( size * sizeof(float));
    m_ArraySize[m_iOutputIndex] = size *sizeof(float);

    // Mat(int rows, int cols, int type, void* data, size_t step=AUTO_STEP);
    // 主要是保存bgr的内存地址
    m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex]);
    m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex] + sizeof(float) * dims_i.d[2] * dims_i.d[3] );
    m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex] + 2 * sizeof(float) * dims_i.d[2] * dims_i.d[3]);
}

（3）对输入进行预处理

这个都是和前面博客讲的差不多，对应yolov5中python源码用c++实现即可（完整代码在下面的github链接）：

	int newh = 0, neww = 0, padh = 0, padw = 0;
	Mat dstimg = this->resize_image(frame, &newh, &neww, &padh, &padw); // resize
	//this->normalize_(dstimg);
	cv::cvtColor(dstimg, dstimg, cv::COLOR_BGR2RGB);   // 由BGR转成RGB
	cv::Mat m_Normalized;
	dstimg.convertTo(m_Normalized, CV_32FC3, 1/255.);  // 这样归一化
    cv::split(m_Normalized, m_InputWrappers);  // 通道分离[h,w,3] ，保存的是rgb顺序

（4）创建CUDA流,推理时TensorRT执行通常是异步的，因此将内核排入CUDA流

	//创建CUDA流,推理时TensorRT执行通常是异步的，因此将内核排入CUDA流
    cudaStreamCreate(&m_CudaStream); 
	auto ret = cudaMemcpyAsync(m_ArrayDevMemory[m_iInputIndex], m_ArrayHostMemory[m_iInputIndex], m_ArraySize[m_iInputIndex], cudaMemcpyHostToDevice, m_CudaStream); 
	auto ret1 = m_CudaContext->enqueueV2(m_ArrayDevMemory, m_CudaStream, nullptr);    // TensorRT 执行通常是异步的，因此将内核排入 CUDA 流：
	ret = cudaMemcpyAsync(m_ArrayHostMemory[m_iOutputIndex], m_ArrayDevMemory[m_iOutputIndex], m_ArraySize[m_iOutputIndex], cudaMemcpyDeviceToHost, m_CudaStream); //输出传回给CPU，数据从显存到内存
    ret = cudaStreamSynchronize(m_CudaStream);

（5）输出结果的生成

	float* pdata = (float*)m_ArrayHostMemory[m_iOutputIndex];
	// 这些就和前面一样，对应yolov5中python源码用c++实现即可：
	std::vector<BoxInfo> generate_boxes;  // BoxInfo自定义的结构体
	float ratioh = (float)frame.rows / newh, ratiow = (float)frame.cols / neww;
	for(int i = 0; i < m_iBoxNums; ++i) // 遍历所有的num_pre_boxes
	{
    
    
		int index = i * (m_iClassNums + 5);      // prob[b*num_pred_boxes*(classes+5)]  
		float obj_conf = pdata[index + 4];  // 置信度分数
		if (obj_conf > this->objThreshold)  // 大于阈值
		{
    
    
			int class_idx = 0;
			float max_class_socre = 0;
			for (int k = 0; k < m_iClassNums; ++k)
			{
    
    
				if (pdata[k + index + 5] > max_class_socre)
				{
    
    
					max_class_socre = pdata[k + index + 5];
					class_idx = k;
				}
			}
			max_class_socre *= obj_conf;   // 最大的类别分数*置信度
			if (max_class_socre > this->confThreshold) // 再次筛选
			{
    
     
				//const int class_idx = classIdPoint.x;
				float cx = pdata[index];  //x
				float cy = pdata[index+1];  //y
				float w = pdata[index+2];  //w
				float h = pdata[index+3];  //h

				float xmin = (cx - padw - 0.5 * w)*ratiow;
				float ymin = (cy - padh - 0.5 * h)*ratioh;
				float xmax = (cx - padw + 0.5 * w)*ratiow;
				float ymax = (cy - padh + 0.5 * h)*ratioh;

				generate_boxes.push_back(BoxInfo{
    
     xmin, ymin, xmax, ymax, max_class_socre, class_idx });
			}
		}
	}

	// Perform non maximum suppression to eliminate redundant overlapping boxes with
	// lower confidences
	nms(generate_boxes);
	for (size_t i = 0; i < generate_boxes.size(); ++i)
	{
    
    
		int xmin = int(generate_boxes[i].x1);
		int ymin = int(generate_boxes[i].y1);
		rectangle(frame, Point(xmin, ymin), Point(int(generate_boxes[i].x2), int(generate_boxes[i].y2)), Scalar(0, 0, 255), 2);
		std::string label = format("%.2f", generate_boxes[i].score);
		label = this->classes[generate_boxes[i].label] + ":" + label;
		putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1);
	}
}

可视化输出

图：
相比于前面opencv里dnn和onnxruntime使用cuda，精度方面几乎没变

时间：
相比于前面opencv里dnn和onnxruntime使用cuda，时间明显要快的多

3.编译

1.直接用c++的编译软件，导入所需头文件和依赖之后直接编译。如：

2.使用cmake编译，添加在了下面的github上，CMakeLists文件里可以方便的用指令导入的头文件和要链接的依赖库，能够灵活的编译不同文件，适合多文件和复杂文件的编译；

总结

c++部署，tensorrt是现在工业上使用c++部署最广的，推理部署的时间很快，还能继续优化，所以有必要掌握部署的整个流程。

整个代码：https://github.com/yzy12-max/yolov5_deploy
参考：
https://docs.nvidia.com/deeplearning/tensorrt/archives/tensorrt-723/developer-guide/index.html
https://docs.nvidia.com/deeplearning/tensorrt/archives/tensorrt-723/api/c_api/
https://github.com/onnx/onnx-tensorrt
https://github.com/NVIDIA/TensorRT
https://github.com/ExtremeMart/ev_sdk_demo4.0_pedestrian_intrusion_yolov5