本文主要是介绍TRT部署ONNX,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
也算是基础知识了 就搬来了 还是要说大佬勿怪 其实应该算是 转成trt默认 使用 这样说确切
不过我一般不这么用yolo直接 用王鑫宇 大佬工具转wts然后载转trt模型 大伙去git上找把 都支持到v7了 是阿拉伯联合酋长国阿布拉卡那版哦
TensorRT
是英伟达官方提供的一个高性能深度学习推理优化库,支持C++
和Python
两种编程语言API
。通常情况下深度学习模型部署都会追求效率,尤其是在嵌入式平台上,所以一般会选择使用C++
来做部署。
这里将以YOLOv5
为例详细介绍如何使用TensorRT
的C++
版本API
来部署ONNX
模型,使用的TensorRT
版本为8.4.1.5
,如果使用其他版本可能会存在某些函数与本文描述的不一致。另外,使用TensorRT 7
会导致YOLOv5
的输出结果与期望不一致,请注意。
导出ONNX模型
YOLOv5
使用PyTorch
框架进行训练,可以使用官方代码仓库中的export.py
脚本把PyTorch
模型转换为ONNX
模型:
python export.py --weights yolov5x.pt --include onnx --imgsz 640 640
准备模型输入数据
如果想用YOLOv5
对图像做目标检测,在将图像输入给模型之前还需要做一定的预处理操作,预处理操作应该与模型训练时所做的操作一致。YOLOv5
的输入是RGB
格式的3
通道图像,图像的每个像素需要除以255
来做归一化,并且数据要按照CHW
的顺序进行排布。所以YOLOv5
的预处理大致可以分为两个步骤:
-
将原始输入图像缩放到模型需要的尺寸,比如
640x640
。这一步需要注意的是,原始图像是按照等比例进行缩放的,如果缩放后的图像某个维度上比目标值小,那么就需要进行填充。举个例子:假设输入图像尺寸为768x576
,模型输入尺寸为640x640
,按照等比例缩放的原则缩放后的图像尺寸为640x480
,那么在y
方向上还需要填充640-480=160
(分别在图像的顶部和底部各填充80
)。来看一下实现代码:
cv::Mat input_image = cv::imread("dog.jpg");
cv::Mat resize_image;
const int model_width = 640;
const int model_height = 640;
const float ratio = std::min(model_width / (input_image.cols * 1.0f), model_height / (input_image.rows * 1.0f));
// 等比例缩放
const int border_width = input_image.cols * ratio;
const int border_height = input_image.rows * ratio;
// 计算偏移值
const int x_offset = (model_width - border_width) / 2;
const int y_offset = (model_height - border_height) / 2;
cv::resize(input_image, resize_image, cv::Size(border_width, border_height));
cv::copyMakeBorder(resize_image, resize_image, y_offset, y_offset, x_offset, x_offset, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));
// 转换为RGB格式
cv::cvtColor(resize_image, resize_image, cv::COLOR_BGR2RGB);
图像这样处理后的效果如下图所示,顶部和底部的灰色部分是填充后的效果。
2. 对图像像素做归一化操作,并按照CHW
的顺序进行排布。这一步的操作比较简单,直接看代码吧:
input_blob = new float[model_height * model_width * 3];
const int channels = resize_image.channels();
const int width = resize_image.cols;
const int height = resize_image.rows;
for (int c = 0; c < channels; c++) { for (int h = 0; h < height; h++) { for (int w = 0; w < width; w++) { input_blob[c * width * height + h * width + w] = resize_image.at<cv::Vec3b>(h, w)[c] / 255.0f; } }
}
ONNX模型部署
1. 模型优化与序列化
要使用TensorRT
的C++ API
来部署模型,首先需要包含头文件NvInfer.h
。
#include "NvInfer.h"
TensorRT
所有的编程接口都被放在命名空间nvinfer1
中,并且都以字母I
为前缀,比如ILogger
、IBuilder
等。使用TensorRT
部署模型首先需要创建一个IBuilder
对象,创建之前还要先实例化ILogger
接口:
class MyLogger : public nvinfer1::ILogger { public: explicit MyLogger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kWARNING) : severity_(severity) {} void log(nvinfer1::ILogger::Severity severity, const char *msg) noexcept override { if (severity <= severity_) { std::cerr << msg << std::endl; } } nvinfer1::ILogger::Severity severity_;
};
上面的代码默认会捕获级别大于等于WARNING
的日志信息并在终端输出。实例化ILogger
接口后,就可以创建IBuilder
对象:
MyLogger logger;
nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);
创建IBuilder
对象后,优化一个模型的第一步是要构建模型的网络结构。
const uint32_t explicit_batch = 1U << static_cast<uint32_t>( nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicit_batch);
模型的网络结构有两种构建方式,一种是使用TensorRT
的API
一层一层地去搭建,这种方式比较麻烦;另外一种是直接从ONNX
模型中解析出模型的网络结构,这需要ONNX
解析器来完成。由于我们已经有现成的ONNX
模型了,所以选择第二种方式。TensorRT
的ONNX
解析器接口被封装在头文件NvOnnxParser.h
中,命名空间为nvonnxparser
。创建ONNX
解析器对象并加载模型的代码如下:
const std::string onnx_model = "yolov5m.onnx";
nvonnxparser::IParser *parser = nvonnxparser::createParser(*network, logger);
parser->parseFromFile(model_path.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kERROR))
// 如果有错误则输出错误信息
for (int32_t i = 0; i < parser->getNbErrors(); ++i) { std::cout << parser->getError(i)->desc() << std::endl;
}
模型解析成功后,需要创建一个IBuilderConfig
对象来告诉TensorRT
该如何对模型进行优化。这个接口定义了很多属性,其中最重要的一个属性是工作空间的最大容量。在网络层实现过程中通常会需要一些临时的工作空间,这个属性会限制最大能申请的工作空间的容量,如果容量不够的话会导致该网络层不能成功实现而导致错误。另外,还可以通过这个对象设置模型的数据精度。TensorRT
默认的数据精度为FP32
,我们还可以设置FP16
或者INT8
,前提是该硬件平台支持这种数据精度。
nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1U << 25);
if (builder->platformHasFastFp16()) { config->setFlag(nvinfer1::BuilderFlag::kFP16);
}
设置IBuilderConfig
属性后,就可以启动优化引擎对模型进行优化了,这个过程需要一定的时间,在嵌入式平台上可能会比较久一点。经过TensorRT
优化后的序列化模型被保存到IHostMemory
对象中,我们可以将其保存到磁盘中,下次使用时直接加载这个经过优化的模型即可,这样就可以省去漫长的等待模型优化的过程。我一般习惯把序列化模型保存到一个后缀为.engine
的文件中。
nvinfer1::IHostMemory *serialized_model = builder->buildSerializedNetwork(*network, *config); // 将模型序列化到engine文件中
std::stringstream engine_file_stream;
engine_file_stream.seekg(0, engine_file_stream.beg);
engine_file_stream.write(static_cast<const char *>(serialized_model->data()), serialized_model->size());
const std::string engine_file_path = "yolov5m.engine";
std::ofstream out_file(engine_file_path);
assert(out_file.is_open());
out_file << engine_file_stream.rdbuf();
out_file.close();
由于IHostMemory
对象保存了模型所有的信息,所以前面创建的IBuilder
、IParser
等对象已经不再需要了,可以通过delete
进行释放。
delete config;
delete parser;
delete network;
delete builder;
IHostMemory
对象用完后也可以通过delete
进行释放。
2. 模型反序列化
通过上一步得到优化后的序列化模型后,如果要用模型进行推理,那么还需要创建一个IRuntime
接口的实例,然后通过其模型反序列化接口去创建一个ICudaEngine
对象:
nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);
nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine( serialized_model->data(), serialized_model->size()); delete serialized_model;
delete runtime;
如果是直接从磁盘中加载.engine
文件也是差不多的步骤,首先从.engine
文件中把模型加载到内存中,然后再通过IRuntime
接口对模型进行反序列化即可。
const std::string engine_file_path = "yolov5m.engine";
std::stringstream engine_file_stream;
engine_file_stream.seekg(0, engine_file_stream.beg);
std::ifstream ifs(engine_file_path);
engine_file_stream << ifs.rdbuf();
ifs.close(); engine_file_stream.seekg(0, std::ios::end);
const int model_size = engine_file_stream.tellg();
engine_file_stream.seekg(0, std::ios::beg);
void *model_mem = malloc(model_size);
engine_file_stream.read(static_cast<char *>(model_mem), model_size); nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);
nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(model_mem, model_size); delete runtime;
free(model_mem);
3. 模型推理
ICudaEngine
对象中存放着经过TensorRT
优化后的模型,不过如果要用模型进行推理则还需要通过createExecutionContext()
函数去创建一个IExecutionContext
对象来管理推理的过程:
nvinfer1::IExecutionContext *context = engine->createExecutionContext();
现在让我们先来看一下使用TensorRT
框架进行模型推理的完整流程:
-
对输入图像数据做与模型训练时一样的预处理操作。
-
把模型的输入数据从
CPU
拷贝到GPU
中。 -
调用模型推理接口进行推理。
-
把模型的输出数据从
GPU
拷贝到CPU
中。 -
对模型的输出结果进行解析,进行必要的后处理后得到最终的结果。
由于模型的推理是在GPU
上进行的,所以会存在搬运输入、输出数据的操作,因此有必要在GPU
上创建内存区域用于存放输入、输出数据。模型输入、输出的尺寸可以通过ICudaEngine
对象的接口来获取,根据这些信息我们可以先为模型分配输入、输出缓存区。
void *buffers[2];
// 获取模型输入尺寸并分配GPU内存
nvinfer1::Dims input_dim = engine->getBindingDimensions(0);
int input_size = 1;
for (int j = 0; j < input_dim.nbDims; ++j) { input_size *= input_dim.d[j];
}
cudaMalloc(&buffers[0], input_size * sizeof(float)); // 获取模型输出尺寸并分配GPU内存
nvinfer1::Dims output_dim = engine->getBindingDimensions(1);
int output_size = 1;
for (int j = 0; j < output_dim.nbDims; ++j) { output_size *= output_dim.d[j];
}
cudaMalloc(&buffers[1], output_size * sizeof(float)); // 给模型输出数据分配相应的CPU内存
float *output_buffer = new float[output_size]();
到这一步,如果你的输入数据已经准备好了,那么就可以调用TensorRT
的接口进行推理了。通常情况下,我们会调用IExecutionContext
对象的enqueueV2()
函数进行异步地推理操作,该函数的第二个参数为CUDA
流对象,第三个参数为CUDA
事件对象,这个事件表示该执行流中输入数据已经使用完,可以挪作他用了。
cudaStream_t stream;
cudaStreamCreate(&stream);
// 拷贝输入数据
cudaMemcpyAsync(buffers[0], input_blob,input_size * sizeof(float), cudaMemcpyHostToDevice, stream);
// 执行推理
context->enqueueV2(buffers, stream, nullptr);
// 拷贝输出数据
cudaMemcpyAsync(output_buffer, buffers[1],output_size * sizeof(float), cudaMemcpyDeviceToHost, stream); cudaStreamSynchronize(stream);
模型推理成功后,其输出数据被拷贝到output_buffer
中,接下来我们只需按照YOLOv5
的输出数据排布规则去解析即可。
4. 小结
在介绍如何解析YOLOv5
输出数据之前,我们先来总结一下用TensorRT
框架部署ONNX
模型的基本流程。
如上图所示,主要步骤如下: whaosoft aiot http://143ai.com
-
实例化
Logger
; -
创建
Builder
; -
创建
Network
; -
使用
Parser
解析ONNX
模型,构建Network
; -
设置
Config
参数; -
优化网络,序列化模型;
-
反序列化模型;
-
拷贝模型输入数据(
HostToDevice
),执行模型推理; -
拷贝模型输出数据(
DeviceToHost
),解析结果。
解析模型输出结果
YOLOv5
有3
个检测头,如果模型输入尺寸为640x640
,那么这3
个检测头分别在80x80
、40x40
和20x20
的特征图上做检测。让我们先用Netron
工具来看一下YOLOv5 ONNX
模型的结构,可以看到,YOLOv5
的后处理操作已经被包含在模型中了(如下图红色框内所示),3
个检测头分支的结果最终被组合成一个张量作为输出。
yolov5m
YOLOv5
的3
个检测头一共有(80x80+40x40+20x20)x3=25200
个输出单元格,每个单元格输出x,y,w,h,objectness
这5
项再加80
个类别的置信度总共85
项内容。经过后处理操作后,目标的坐标值已经被恢复到以640x640
为参考的尺寸,如果需要恢复到原始图像尺寸,只需要除以预处理时的缩放因子即可。这里有个问题需要注意:由于在做预处理的时候图像做了填充,原始图像并不是被缩放成640x640
而是640x480
,使得输入给模型的图像的顶部被填充了一块高度为80
的区域,所以在恢复到原始尺寸之前,需要把目标的y
坐标减去偏移量80
。
详细的解析代码如下:
float *ptr = output_buffer;
for (int i = 0; i < 25200; ++i) { const float objectness = ptr[4]; if (objectness >= 0.45f) { const int label = std::max_element(ptr + 5, ptr + 85) - (ptr + 5); const float confidence = ptr[5 + label] * objectness; if (confidence >= 0.25f) { const float bx = ptr[0]; const float by = ptr[1]; const float bw = ptr[2]; const float bh = ptr[3]; Object obj; // 这里要减掉偏移值 obj.box.x = (bx - bw * 0.5f - x_offset) / ratio; obj.box.y = (by - bh * 0.5f - y_offset) / ratio; obj.box.width = bw / ratio; obj.box.height = bh / ratio; obj.label = label; obj.confidence = confidence; objs->push_back(std::move(obj)); } } ptr += 85;
} // i loop
对解析出的目标做非极大值抑制(NMS
)操作后,检测结果如下图所示:
本文以YOLOv5
为例通过大量的代码一步步讲解如何使用TensorRT
框架部署ONNX
模型,主要目的是希望读者能够通过本文学习到TensorRT
模型部署的基本流程,比如如何准备输入数据、如何调用API
用模型做推理、如何解析模型的输出结果。如何部署YOLOv5
模型并不是本文的重点,重点是要掌握使用TensorRT
部署ONNX
模型的基本方法,这样才会有举一反三的效果。
这篇关于TRT部署ONNX的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!