YOLOV5在C++中通过TensorRT推理
目录
项目地址：https://gitee.com/lazyone/InferOnnx.git
需要提前下好spdlog：https://github.com/gabime/spdlog
完整资源：YOLOv5+TensorRT/OnnxRuntime+Visual Studio+CmakeLists实现推理
Onnx推理文章：Onnxruntime在window上的cpu与gpu推理
环境设置
CPU:i5-9400F
GPU:GTX1060
TensorRT:8.5.1
OnnxRuntime:1.13
参考文档
yolov5使用TensorRT进行c++部署:跳转链接
NVIDIA TENSORRT DOCUMENTATION:跳转链接
onnx-tensorrt:跳转链接
极市YOLOV5的TensorRT部署Demo：跳转链接
推理结果
直接看onnxruntime与tensorRT在FP32的情况下加速推理的结果比较
可以看到tensorRT还是有很大的优势的。
对比下输出结果，首先是tensorRT的输出结果：
接着是OnnxRuntime的输出结果：
可以看到结果基本一样。
代码分析
主要分析和TensorRT有关的，其余的前处理和后处理都大同小异。
1. 将ONNX模型转换为TensorRT的Engine文件
如果已经有了TensorRT的Engine文件，可以直接跳过这一步。转换方式这里介绍两种：
-
使用TensorRT提供的
trtexec工具将其转换为Engine文件。例如，假设ONNX文件命名为model.onnx，执行以下命令可以将其转换为Engine文件：复制代码
trtexec --onnx=model.onnx --explicitBatch上述命令中，
--explicitBatch表示明确指定输入数据的batch size，这样在执行推理时可以动态地调整batch size的大小。 -
根据官方代码转换成trt文件并加载
void YOLOv5::loadOnnx(const std::string strModelName) { TRT_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(); std::cout << "convert onnx model to TensorRT engine model successfully!" << std::endl; }TRT_Logger类是nvinfer1::ILogger的子类nvinfer1::ILogger是TensorRT中定义的日志输出接口，用于输出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; } } };TRT_Logger类中定义了log函数，用于输出日志信息。它接受两个参数：severity：表示日志信息的级别，是一个nvinfer1::ILogger::Severity类型的变量。msg：表示日志信息的字符串。
2. 加载Engine文件
如果是读取Engine文件，则需要加载Engine。首先创建了一个nvinfer1::IRuntime的对象，然后加载指定路径下的Engine文件（其路径为enginePath），并使用反序列化方法deserializeCudaEngine将其反序列化。接着，创建了一个m_CudaContext执行上下文对象，用于执行推理。
void YOLOv5::loadTrt(const std::string strName) {
TRT_Logger gLogger;
// 创建 TensorRT 运行时对象
m_CudaRuntime = nvinfer1::createInferRuntime(gLogger);
// 打开引擎文件，读取序列化数据
std::ifstream fin(strName, std::ios::binary);
if (!fin.good()) {
fin.close();
throw std::runtime_error("Unable to open engine file: " + strName);
}
// 获取文件大小
fin.seekg(0, std::ios::end);
size_t fileSize = fin.tellg();
fin.seekg(0, std::ios::beg);
// 分配缓冲区，存储序列化数据
std::unique_ptr<char[]> engineBuffer(new char[fileSize]);
fin.read(engineBuffer.get(), fileSize);
fin.close();
// 反序列化引擎
m_CudaEngine = m_CudaRuntime->deserializeCudaEngine(engineBuffer.get(),
fileSize, nullptr);
if (!m_CudaEngine) {
throw std::runtime_error("Failed to deserialize engine from file: " +
strName);
}
// 创建 TensorRT 执行上下文
m_CudaContext = m_CudaEngine->createExecutionContext();
// 释放内存
m_CudaRuntime->destroy();
}
3. 对输入数据进行处理
首先需要获取输入与输出节点的相关信息来计算设备中需要申请的内存大小
// 分配内存大小
cudaMalloc(
&m_ArrayDevMemory[m_iInputIndex],
size1 *
sizeof(
float)); // 在CUDA显存上分配size1 * sizeof(float)大小的内存空间，并将该显存指针赋值给m_ArrayDevMemory[m_iInputIndex]
cudaMalloc(&m_ArrayDevMemory[m_iOutputIndex], size2 * sizeof(float)); // 同上
m_ArrayHostMemory[m_iInputIndex] = malloc(
size1 *
sizeof(
float)); // 于在主机内存上分配一块大小为size1 * sizeof(float)的内存空间，并将该地址赋值给m_ArrayHostMemory[m_iInputIndex]
m_ArrayHostMemory[m_iOutputIndex] = malloc(size2 * sizeof(float)); //
m_ArraySize[m_iInputIndex] =
size1 *
sizeof(
float); // 将当前输入数据的内存空间大小记录在m_ArraySize数组中的对应位置，以方便后续的内存传输操作。
m_ArraySize[m_iOutputIndex] = size2 * sizeof(float); //
将TensorRT推理的输入数据转换为OpenCV中的Mat格式，m_InputWrappers是个vector向量，包含三个Mat，存储了三个通道的数据：
m_InputWrappers.emplace_back(
dims_i.d[2], dims_i.d[3], CV_32FC1,
m_ArrayHostMemory
[m_iInputIndex]); // 将m_ArrayHostMemory[m_iInputIndex]指向的一段空间按照每个元素大小为一个float的方式解析成一个矩阵（Mat），其维度为dims_i.d[2]×dims_i.d[3]，像素类型是CV_32FC1（32位浮点数单通道）。
m_InputWrappers.emplace_back(
dims_i.d[2], dims_i.d[3], CV_32FC1,
(char*)m_ArrayHostMemory[m_iInputIndex] +
sizeof(float) * dims_i.d[2] *
dims_i.d
[3]); // visual studio 不支持对void*的指针直接运算，需要转换成具体类型。，
m_InputWrappers.emplace_back(
dims_i.d[2], dims_i.d[3], CV_32FC1,
(char*)m_ArrayHostMemory[m_iInputIndex] +
2 * sizeof(float) * dims_i.d[2] * dims_i.d[3]);
需要注意到，Mat数据data地址已经被指定好了。
4. 执行推理
在输入数据处理完成后，就可以使用执行上下文对象进行推理。先进行相关预处理
cv::cvtColor(dstimg, dstimg, cv::COLOR_BGR2RGB); // 由BGR转成RGB
cv::Mat m_Normalized;
dstimg.convertTo(m_Normalized, CV_32FC3, 1 / 255.);
// 将m_Normalized的每个通道拆分为一个矩阵，然后将这些矩阵按顺序存储到m_InputWrappers的数组中，以替换现有内容。
cv::split(m_Normalized, m_InputWrappers); // 通道分离[h,w,3] RGB
注意图像通道的顺序，将数据存到指定地址。
// 创建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);
float* pdata = (float*)m_ArrayHostMemory[m_iOutputIndex];
-
创建CUDA流，用于异步执行内核操作。
cudaStreamCreate(&m_CudaStream); -
使用异步
cudaMemcpy将输入数据从主机内存拷贝到设备内存。auto ret = cudaMemcpyAsync( m_ArrayDevMemory[m_iInputIndex], m_ArrayHostMemory[m_iInputIndex], m_ArraySize[m_iInputIndex], cudaMemcpyHostToDevice, m_CudaStream); -
将内核排入CUDA流，异步执行TensorRT推理。
auto ret1 = m_CudaContext->enqueueV2( m_ArrayDevMemory, m_CudaStream, nullptr); -
使用异步
cudaMemcpy将输出数据从设备内存拷贝到主机内存。ret = cudaMemcpyAsync(m_ArrayHostMemory[m_iOutputIndex], m_ArrayDevMemory[m_iOutputIndex], m_ArraySize[m_iOutputIndex], cudaMemcpyDeviceToHost, m_CudaStream); -
同步CUDA流，等待所有处理任务完成。
ret = cudaStreamSynchronize(m_CudaStream); -
将
m_ArrayHostMemory[m_iOutputIndex]的数据转换为浮点型指针。float* pdata = (float*)m_ArrayHostMemory[m_iOutputIndex];
整个过程中，推理和数据传输的过程均被异步化，可以提高推理速度和资源利用率。注意，此处的CUDA流只是用于异步执行内核操作，并非用于并行操作。这里需要一定的cuda知识（我也不懂，大家都这样写，以后在研究）可以参考：
https://developer.nvidia.com/zh-cn/blog/tensorrt-measuring-performance-cn/
5. 释放资源
void YOLOv5::UnInit() {
for (auto& p : m_ArrayDevMemory) {
cudaFree(p);
p = nullptr;
}
for (auto& p : m_ArrayHostMemory) {
free(p);
p = nullptr;
}
cudaStreamDestroy(m_CudaStream);
m_CudaContext->destroy(); // 首先释放该CUDA上下文使用的CUDA资源
m_CudaEngine->destroy(); // 再释放该TensorRT模型的CUDA引擎对象。
}
注意下释放顺序，否则会报错。
知识点
- emplace_back()方法可以直接在vector尾部构造元素，而不是创建临时对象，然后将其拷贝到vector中。这样可以避免调用拷贝构造函数和移动构造函数，提高效率。
- 矩阵的构造：cv::Mat(int rows, int cols, int type, void* data, size_t step = AUTO_STEP);
- 执行下段语句时
由于visual studio 不支持对m_InputWrappers.emplace_back( dims_i.d[2], dims_i.d[3], CV_32FC1, (char*)m_ArrayHostMemory[m_iInputIndex] + sizeof(float) * dims_i.d[2] * dims_i.d [3]);void*的指针直接运算，需要转换成具体类型。这里转成float*会导致程序出错，具体原因不太明白。搞了好久，转成char*就可以了。 - 文件流操作
该段代码中使用了// 使用std::ifstream打开引擎文件，读入二进制数据。 std::ifstream fin(strName, std::ios::binary); // 检查文件是否成功打开，如果没有成功打开，则抛出运行时异常。 if (!fin.good()) { fin.close(); throw std::runtime_error("Unable to open engine file: " + strName); } // 获取文件大小。 fin.seekg(0, std::ios::end); size_t fileSize = fin.tellg(); fin.seekg(0, std::ios::beg); // 分配缓冲区，存储序列化数据。 /* 首先通过 new 运算符在堆上动态分配一个 char 数组，该数组可以容纳文件中的所有数据。 new char[fileSize] 然后把分配得到的数组转换成一个 std::unique_ptr 智能指针；这个智能指针是用来防止内存泄漏的，并在其生命周期内自动管理所分配的数组内存。 std::unique_ptr<char[]> engineBuffer(new char[fileSize]) 接下来通过 std::ifstream 的 read 成员函数，从文件流中读取全部 fileSize 个字节的二进制数据，并写入到 engineBuffer 所指向的缓冲区中。其中，get() 得到所管理的原始指针。 fin.read(engineBuffer.get(), fileSize) 最终，这段代码能够将整个二进制文件的数据流读到内存中，并以 char 数组的形式存储在engineBuffer所管理的内存区域中。 */ std::unique_ptr<char[]> engineBuffer(new char[fileSize]); fin.read(engineBuffer.get(), fileSize); // 关闭文件流。 fin.close();std::ifstream读入引擎文件二进制数据，并在内存中分配了一个缓冲区，用于存储序列化的引擎数据。这个缓冲区后续会被用于创建TensorRT引擎。此外，还使用std::ios::binary表示以二进制方式打开文件，这也是读取二进制数据的常用方式。
代码
cpp文件
#include "tensorRT.h"
#include <sys/stat.h>
#include <ctime>
#include <fstream>
#include <iostream>
// 命名空间
using namespace cv;
using namespace nvinfer1;
// 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;
}
}
};
// 判断文件是否形成
static bool ifFileExists(const char* FileName) {
struct stat my_stat;
return (stat(FileName, &my_stat) == 0);
}
// 加载onnx文件
void YOLOv5::loadOnnx(const std::string strModelName) {
TRT_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();
std::cout << "convert onnx model to TensorRT engine model successfully!"
<< std::endl;
}
/**
*引入 nvinfer1 命名空间，以访问 TensorRT 8.5 版本的 API。
在打开引擎文件的时候需要指定 std::ios::binary 参数，以保证二进制文件的读取和写入方式。
在序列化文件的读取方法上，使用了一个 std::stringstream 来承载 std::ifstream::rdbuf() 或者 std::istreambuf_iterator<char>(fin) 的全部读入，以便于观察序列化数据的类型。
在反序列化引擎时，将引擎序列化数据存储在缓冲区中，并使用大小安全的缓冲区指针和大小传递给函数。
在创建 TensorRT 执行上下文时，先判断引擎是否为空，避免出现程序崩溃等问题。
需要注意的是，在修改和优化代码时，应该进行充分测试和验证，确保代码的正确性和健壮性。
同时，需要注意 TensorRT 版本的兼容性，确保在不同版本的 TensorRT 上都能够正确运行代码。
*/
void YOLOv5::loadTrt(const std::string strName) {
TRT_Logger gLogger;
// 创建 TensorRT 运行时对象
m_CudaRuntime = nvinfer1::createInferRuntime(gLogger);
// 打开引擎文件，读取序列化数据
std::ifstream fin(strName, std::ios::binary);
if (!fin.good()) {
fin.close();
throw std::runtime_error("Unable to open engine file: " + strName);
}
// 获取文件大小
fin.seekg(0, std::ios::end);
size_t fileSize = fin.tellg();
fin.seekg(0, std::ios::beg);
// 分配缓冲区，存储序列化数据
std::unique_ptr<char[]> engineBuffer(new char[fileSize]);
fin.read(engineBuffer.get(), fileSize);
fin.close();
// 反序列化引擎
m_CudaEngine = m_CudaRuntime->deserializeCudaEngine(engineBuffer.get(),
fileSize, nullptr);
if (!m_CudaEngine) {
throw std::runtime_error("Failed to deserialize engine from file: " +
strName);
}
// 创建 TensorRT 执行上下文
m_CudaContext = m_CudaEngine->createExecutionContext();
// 释放内存
m_CudaRuntime->destroy();
}
// 初始化
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"; //
if (ifFileExists(strTrtName.c_str())) {
loadTrt(strTrtName);
} else {
loadOnnx(config.modelpath);
}
// 利用加载的模型获取输入输出信息
// 使用输入和输出blob名来获取输入和输出索引
m_iInputIndex = m_CudaEngine->getBindingIndex("images"); // 输入索引
m_iOutputIndex = m_CudaEngine->getBindingIndex("output0"); // 输出
// 输出处理数据类型
std::cout << "type = ";
switch (m_CudaEngine->getBindingDataType(m_iInputIndex)) {
case nvinfer1::DataType::kFLOAT:
std::cout << "FP32";
break;
case nvinfer1::DataType::kHALF:
std::cout << "FP16";
break;
case nvinfer1::DataType::kINT8:
std::cout << "INT8";
break;
case nvinfer1::DataType::kINT32:
std::cout << "INT32";
break;
default:
std::cout << "unknown";
break;
}
std::cout << std::endl;
Dims dims_i = m_CudaEngine->getBindingDimensions(
m_iInputIndex); // 输入，维度[0,1,2,3]NHWC
Dims dims_o = m_CudaEngine->getBindingDimensions(m_iOutputIndex); // 输出
std::cout << "input dims： " << dims_i.nbDims << " " << dims_i.d[0] << " "
<< dims_i.d[1] << " " << dims_i.d[2] << " " << dims_i.d[3] << endl;
std::cout << "output dims： " << dims_o.nbDims << " " << dims_o.d[0] << " "
<< dims_o.d[1] << " " << dims_o.d[2] << endl;
int size1 = dims_i.d[0] * dims_i.d[1] * dims_i.d[2] * dims_i.d[3]; // 展平
int size2 = dims_o.d[0] * dims_o.d[1] * dims_o.d[2]; // 所有大小
m_InputSize = cv::Size(dims_i.d[3], dims_i.d[2]); // 输入尺寸(W,H)
m_iClassNums = dims_o.d[2] - 5; // 类别数量
m_iBoxNums = dims_o.d[1]; // num_pre_boxes
// 分配内存大小
cudaMalloc(
&m_ArrayDevMemory[m_iInputIndex],
size1 *
sizeof(
float)); // 在CUDA显存上分配size1 * sizeof(float)大小的内存空间，并将该显存指针赋值给m_ArrayDevMemory[m_iInputIndex]
cudaMalloc(&m_ArrayDevMemory[m_iOutputIndex], size2 * sizeof(float)); // 同上
m_ArrayHostMemory[m_iInputIndex] = malloc(
size1 *
sizeof(
float)); // 于在主机内存上分配一块大小为size1 * sizeof(float)的内存空间，并将该地址赋值给m_ArrayHostMemory[m_iInputIndex]
m_ArrayHostMemory[m_iOutputIndex] = malloc(size2 * sizeof(float)); //
m_ArraySize[m_iInputIndex] =
size1 *
sizeof(
float); // 将当前输入数据的内存空间大小记录在m_ArraySize数组中的对应位置，以方便后续的内存传输操作。
m_ArraySize[m_iOutputIndex] = size2 * sizeof(float); //
// 知识点：
// 1. emplace_back()方法可以直接在vector尾部构造元素，而不是创建临时对象，然后将其拷贝到vector中。这样可以避免调用拷贝构造函数和移动构造函数，提高效率。
// 2. 矩阵的构造：cv::Mat(int rows, int cols, int type, void* data, size_t step = AUTO_STEP);
m_InputWrappers.emplace_back(
dims_i.d[2], dims_i.d[3], CV_32FC1,
m_ArrayHostMemory
[m_iInputIndex]); // 将m_ArrayHostMemory[m_iInputIndex]指向的一段空间按照每个元素大小为一个float的方式解析成一个矩阵（Mat），其维度为dims_i.d[2]×dims_i.d[3]，像素类型是CV_32FC1（32位浮点数单通道）。
m_InputWrappers.emplace_back(
dims_i.d[2], dims_i.d[3], CV_32FC1,
(char*)m_ArrayHostMemory[m_iInputIndex] +
sizeof(float) * dims_i.d[2] *
dims_i.d
[3]); // visual studio 不支持对void*的指针直接运算，需要转换成具体类型。，
m_InputWrappers.emplace_back(
dims_i.d[2], dims_i.d[3], CV_32FC1,
(char*)m_ArrayHostMemory[m_iInputIndex] +
2 * sizeof(float) * dims_i.d[2] * dims_i.d[3]);
// 获取类别信息
std::string classesFile =
"D:/workspace/C++/Onnx/InferOnxx/InferOnxx/class.names";
std::ifstream ifs(classesFile.c_str());
std::string line;
while (getline(ifs, line)) this->class_names.push_back(line);
std::cout << "Finish Init TensorRT" << endl;
}
void YOLOv5::UnInit() {
for (auto& p : m_ArrayDevMemory) {
cudaFree(p);
p = nullptr;
}
for (auto& p : m_ArrayHostMemory) {
free(p);
p = nullptr;
}
cudaStreamDestroy(m_CudaStream);
m_CudaContext->destroy(); // 首先释放该CUDA上下文使用的CUDA资源
m_CudaEngine->destroy(); // 再释放该TensorRT模型的CUDA引擎对象。
}
YOLOv5::~YOLOv5() { UnInit(); }
Mat YOLOv5::resize_image(Mat srcimg, int* newh, int* neww, int* top,
int* left) {
int srch = srcimg.rows, srcw = srcimg.cols;
*newh = this->inpHeight;
*neww = this->inpWidth;
Mat dstimg;
if (this->keep_ratio && srch != srcw) {
float hw_scale = (float)srch / srcw;
if (hw_scale > 1) {
*newh = this->inpHeight;
*neww = int(this->inpWidth / hw_scale);
resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
*left = int((this->inpWidth - *neww) * 0.5);
copyMakeBorder(dstimg, dstimg, 0, 0, *left,
this->inpWidth - *neww - *left, BORDER_CONSTANT,
cv::Scalar(114, 114, 114));
} else {
*newh = (int)this->inpHeight * hw_scale;
*neww = this->inpWidth;
resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
*top = (int)(this->inpHeight - *newh) * 0.5;
copyMakeBorder(dstimg, dstimg, *top, this->inpHeight - *newh - *top, 0, 0,
BORDER_CONSTANT, cv::Scalar(114, 114, 114));
}
} else {
resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
}
return dstimg;
}
void YOLOv5::nms(std::vector<BoxInfo>& input_boxes) {
sort(input_boxes.begin(), input_boxes.end(),
[](BoxInfo a, BoxInfo b) { return a.score > b.score; }); // 降序排列
std::vector<bool> remove_flags(input_boxes.size(), false);
auto iou = [](const BoxInfo& box1, const BoxInfo& box2) {
float xx1 = max(box1.x1, box2.x1);
float yy1 = max(box1.y1, box2.y1);
float xx2 = min(box1.x2, box2.x2);
float yy2 = min(box1.y2, box2.y2);
// 交集
float w = max(0.0f, xx2 - xx1 + 1);
float h = max(0.0f, yy2 - yy1 + 1);
float inter_area = w * h;
// 并集
float union_area =
max(0.0f, box1.x2 - box1.x1) * max(0.0f, box1.y2 - box1.y1) +
max(0.0f, box2.x2 - box2.x1) * max(0.0f, box2.y2 - box2.y1) -
inter_area;
return inter_area / union_area;
};
for (int i = 0; i < input_boxes.size(); ++i) {
if (remove_flags[i]) continue;
for (int j = i + 1; j < input_boxes.size(); ++j) {
if (remove_flags[j]) continue;
if (input_boxes[i].label == input_boxes[j].label &&
iou(input_boxes[i], input_boxes[j]) >= this->nmsThreshold) {
remove_flags[j] = true;
}
}
}
int idx_t = 0;
// remove_if()函数 remove_if(beg, end, op) // 移除区间[beg,end)中每一个“令判断式:op(elem)获得true”的元素
input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(),
[&idx_t, &remove_flags](const BoxInfo& f) {
return remove_flags[idx_t++];
}),
input_boxes.end());
}
void YOLOv5::detect(Mat& frame) {
int newh = 0, neww = 0, padh = 0, padw = 0;
Mat dstimg = this->resize_image(frame, &newh, &neww, &padh, &padw);
cv::cvtColor(dstimg, dstimg, cv::COLOR_BGR2RGB); // 由BGR转成RGB
cv::Mat m_Normalized;
dstimg.convertTo(m_Normalized, CV_32FC3, 1 / 255.);
// 将m_Normalized的每个通道拆分为一个矩阵，然后将这些矩阵按顺序存储到m_InputWrappers的数组中，以替换现有内容。
cv::split(m_Normalized, m_InputWrappers); // 通道分离[h,w,3] RGB
// 创建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);
float* pdata = (float*)m_ArrayHostMemory[m_iOutputIndex];
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) // 大于阈值
{
float* max_class_pos = std::max_element(
pdata + index + 5, pdata + index + 5 + m_iClassNums); //
(*max_class_pos) *= obj_conf; // 最大的类别分数*置信度
if ((*max_class_pos) > 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_pos),
static_cast<int>(max_class_pos - (pdata + index + 5))});
}
}
}
// 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(255, 0, 255), 2);
std::string label = format("%.2f", generate_boxes[i].score);
label = this->class_names[generate_boxes[i].label] + ":" + label;
putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75,
Scalar(0, 255, 255), 2);
}
}
h文件
#ifndef __TENSORRT_H__
#define __TENSORRT_H__
#include <NvInfer.h> // nvidia加载模型进行推理的插件
#include <NvOnnxParser.h>
#include <cuda_runtime.h>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include "InferOnxx.h"
// 自定义配置结构
struct Configuration {
float confThreshold; // Confidence threshold
float nmsThreshold; // Non-maximum suppression threshold
float objThreshold; // Object Confidence threshold
std::string modelpath;
};
class YOLOv5 {
public:
YOLOv5(Configuration config);
~YOLOv5();
void UnInit();
void detect(cv::Mat &frame);
private:
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);
void loadTrt(const std::string strName);
private:
float confThreshold;
float nmsThreshold;
float objThreshold;
int inpWidth;
int inpHeight;
std::vector<std::string> class_names; // 类别名称
const bool keep_ratio = true;
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};
void *m_ArrayHostMemory[2]{0};
int m_ArraySize[2]{0};
std::vector<cv::Mat> m_InputWrappers{}; // 空的matvector
};
#endif //__TENSORRT_H__
end
扫一扫
5915
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
