0
  • 聊天消息
  • 系统消息
  • 评论与回复
登录后你可以
  • 下载海量资料
  • 学习在线课程
  • 观看技术视频
  • 写文章/发帖/加入社区
会员中心
创作中心

完善资料让更多小伙伴认识你,还能领取20积分哦,立即完善>

3天内不再提示

如何使用TensorRT框架部署ONNX模型

新机器视觉 来源:新机器视觉 作者:新机器视觉 2022-10-31 14:27 次阅读

导读 模型部署作为算法模型落地的最后一步,在人工智能产业化过程中是非常关键的步骤,而目标检测作为计算机视觉三大基础任务之一,众多的业务功能都要在检测的基础之上完成,本文提供了YOLOv5算法从0部署的实战教程,值得各位读者收藏学习。

前言

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模型:


pythonexport.py--weightsyolov5x.pt--includeonnx--imgsz640640

准备模型输入数据

如果想用YOLOv5对图像做目标检测,在将图像输入给模型之前还需要做一定的预处理操作,预处理操作应该与模型训练时所做的操作一致。YOLOv5的输入是RGB格式的3通道图像,图像的每个像素需要除以255来做归一化,并且数据要按照CHW的顺序进行排布。所以YOLOv5的预处理大致可以分为两个步骤:

将原始输入图像缩放到模型需要的尺寸,比如640x640。这一步需要注意的是,原始图像是按照等比例进行缩放的,如果缩放后的图像某个维度上比目标值小,那么就需要进行填充。举个例子:假设输入图像尺寸为768x576,模型输入尺寸为640x640,按照等比例缩放的原则缩放后的图像尺寸为640x480,那么在y方向上还需要填充640-480=160(分别在图像的顶部和底部各填充80)。来看一下实现代码:


cv::Matinput_image=cv::imread("dog.jpg"); cv::Matresize_image; constintmodel_width=640; constintmodel_height=640; constfloatratio=std::min(model_width/(input_image.cols*1.0f), model_height/(input_image.rows*1.0f)); //等比例缩放 constintborder_width=input_image.cols*ratio; constintborder_height=input_image.rows*ratio; //计算偏移值 constintx_offset=(model_width-border_width)/2; constinty_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); 图像这样处理后的效果如下图所示,顶部和底部的灰色部分是填充后的效果。

对图像像素做归一化操作,并按照CHW的顺序进行排布。这一步的操作比较简单,直接看代码吧:


input_blob=newfloat[model_height*model_width*3]; constintchannels=resize_image.channels(); constintwidth=resize_image.cols; constintheight=resize_image.rows; for(intc=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(h,w)[c]/255.0f; } } }

ONNX模型部署

1. 模型优化与序列化

要使用TensorRT的C++ API来部署模型,首先需要包含头文件NvInfer.h。


#include"NvInfer.h" TensorRT所有的编程接口都被放在命名空间nvinfer1中,并且都以字母I为前缀,比如ILogger、IBuilder等。使用TensorRT部署模型首先需要创建一个IBuilder对象,创建之前还要先实例化ILogger接口:

classMyLogger:publicnvinfer1::ILogger{ public: explicitMyLogger(nvinfer1::Severityseverity= nvinfer1::kWARNING) :severity_(severity){} voidlog(nvinfer1::Severityseverity, constchar*msg)noexceptoverride{ if(severity<= severity_) {       std::cerr << msg << std::endl;     }   }   nvinfer1::Severity severity_; }; 上面的代码默认会捕获级别大于等于WARNING的日志信息并在终端输出。实例化ILogger接口后,就可以创建IBuilder对象:

MyLoggerlogger; nvinfer1::IBuilder*builder=nvinfer1::createInferBuilder(logger); 创建IBuilder对象后,优化一个模型的第一步是要构建模型的网络结构。

constuint32_texplicit_batch=1U<< static_cast( nvinfer1::kEXPLICIT_BATCH); nvinfer1::INetworkDefinition*network=builder->createNetworkV2(explicit_batch); 模型的网络结构有两种构建方式,一种是使用TensorRT的API一层一层地去搭建,这种方式比较麻烦;另外一种是直接从ONNX模型中解析出模型的网络结构,这需要ONNX解析器来完成。由于我们已经有现成的ONNX模型了,所以选择第二种方式。TensorRT的ONNX解析器接口被封装在头文件NvOnnxParser.h中,命名空间为nvonnxparser。创建ONNX解析器对象并加载模型的代码如下:

conststd::stringonnx_model="yolov5m.onnx"; nvonnxparser::IParser*parser=nvonnxparser::createParser(*network,logger); parser->parseFromFile(model_path.c_str(), static_cast(nvinfer1::kERROR)) //如果有错误则输出错误信息 for(int32_ti=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::kWORKSPACE,1U<< 25); if (builder->platformHasFastFp16()){ config->setFlag(nvinfer1::kFP16); } 设置IBuilderConfig属性后,就可以启动优化引擎对模型进行优化了,这个过程需要一定的时间,在嵌入式平台上可能会比较久一点。经过TensorRT优化后的序列化模型被保存到IHostMemory对象中,我们可以将其保存到磁盘中,下次使用时直接加载这个经过优化的模型即可,这样就可以省去漫长的等待模型优化的过程。我一般习惯把序列化模型保存到一个后缀为.engine的文件中。

nvinfer1::IHostMemory*serialized_model= builder->buildSerializedNetwork(*network,*config); //将模型序列化到engine文件中 std::stringstreamengine_file_stream; engine_file_stream.seekg(0,engine_file_stream.beg); engine_file_stream.write(static_cast(serialized_model->data()), serialized_model->size()); conststd::stringengine_file_path="yolov5m.engine"; std::ofstreamout_file(engine_file_path); assert(out_file.is_open()); out_file<< engine_file_stream.rdbuf(); out_file.close(); 由于IHostMemory对象保存了模型所有的信息,所以前面创建的IBuilder、IParser等对象已经不再需要了,可以通过delete进行释放。

deleteconfig; deleteparser; deletenetwork; deletebuilder; IHostMemory对象用完后也可以通过delete进行释放。

2. 模型反序列化

通过上一步得到优化后的序列化模型后,如果要用模型进行推理,那么还需要创建一个IRuntime接口的实例,然后通过其模型反序列化接口去创建一个ICudaEngine对象:


nvinfer1::IRuntime*runtime=nvinfer1::createInferRuntime(logger); nvinfer1::ICudaEngine*engine=runtime->deserializeCudaEngine( serialized_model->data(),serialized_model->size()); deleteserialized_model; deleteruntime; 如果是直接从磁盘中加载.engine文件也是差不多的步骤,首先从.engine文件中把模型加载到内存中,然后再通过IRuntime接口对模型进行反序列化即可。

conststd::stringengine_file_path="yolov5m.engine"; std::stringstreamengine_file_stream; engine_file_stream.seekg(0,engine_file_stream.beg); std::ifstreamifs(engine_file_path); engine_file_stream<< ifs.rdbuf(); ifs.close(); engine_file_stream.seekg(0, std::end); const int model_size = engine_file_stream.tellg(); engine_file_stream.seekg(0, std::beg); void *model_mem = malloc(model_size); engine_file_stream.read(static_cast(model_mem),model_size); nvinfer1::IRuntime*runtime=nvinfer1::createInferRuntime(logger); nvinfer1::ICudaEngine*engine=runtime->deserializeCudaEngine(model_mem,model_size); deleteruntime; free(model_mem);

3. 模型推理

ICudaEngine对象中存放着经过TensorRT优化后的模型,不过如果要用模型进行推理则还需要通过createExecutionContext()函数去创建一个IExecutionContext对象来管理推理的过程:


nvinfer1::IExecutionContext*context=engine->createExecutionContext(); 现在让我们先来看一下使用TensorRT框架进行模型推理的完整流程: 4a528b64-58d6-11ed-a3b6-dac502259ad0.png

对输入图像数据做与模型训练时一样的预处理操作。

把模型的输入数据从CPU拷贝到GPU中。

调用模型推理接口进行推理。

把模型的输出数据从GPU拷贝到CPU中。

对模型的输出结果进行解析,进行必要的后处理后得到最终的结果。

由于模型的推理是在GPU上进行的,所以会存在搬运输入、输出数据的操作,因此有必要在GPU上创建内存区域用于存放输入、输出数据。模型输入、输出的尺寸可以通过ICudaEngine对象的接口来获取,根据这些信息我们可以先为模型分配输入、输出缓存区。


void*buffers[2]; //获取模型输入尺寸并分配GPU内存 nvinfer1::Dimsinput_dim=engine->getBindingDimensions(0); intinput_size=1; for(intj=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); intoutput_size=1; for(intj=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事件对象,这个事件表示该执行流中输入数据已经使用完,可以挪作他用了。如果对CUDA的流和事件不了解,可以参考我之前写的这篇文章。

cudaStream_tstream; 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模型的基本流程。 4a59523c-58d6-11ed-a3b6-dac502259ad0.png 如上图所示,主要步骤如下:

实例化Logger;

创建Builder;

创建Network;

使用Parser解析ONNX模型,构建Network;

设置Config参数

优化网络,序列化模型;

反序列化模型;

拷贝模型输入数据(HostToDevice),执行模型推理;

拷贝模型输出数据(DeviceToHost),解析结果。

解析模型输出结果

YOLOv5有3个检测头,如果模型输入尺寸为640x640,那么这3个检测头分别在80x80、40x40和20x20的特征图上做检测。让我们先用Netron工具来看一下YOLOv5 ONNX模型的结构,可以看到,YOLOv5的后处理操作已经被包含在模型中了(如下图红色框内所示),3个检测头分支的结果最终被组合成一个张量作为输出。 4a73c8d8-58d6-11ed-a3b6-dac502259ad0.pngyolov5m YOLOv5的3个检测头一共有(80x80+40x40+20x20)x3=25200个输出单元格,每个单元格输出x,y,w,h,objectness这5项再加80个类别的置信度总共85项内容。经过后处理操作后,目标的坐标值已经被恢复到以640x640为参考的尺寸,如果需要恢复到原始图像尺寸,只需要除以预处理时的缩放因子即可。这里有个问题需要注意:由于在做预处理的时候图像做了填充,原始图像并不是被缩放成640x640而是640x480,使得输入给模型的图像的顶部被填充了一块高度为80的区域,所以在恢复到原始尺寸之前,需要把目标的y坐标减去偏移量80。 4a7e7396-58d6-11ed-a3b6-dac502259ad0.png 详细的解析代码如下:


float*ptr=output_buffer; for(inti=0;i< 25200; ++i) {   const float objectness = ptr[4];   if (objectness >=0.45f){ constintlabel= std::max_element(ptr+5,ptr+85)-(ptr+5); constfloatconfidence=ptr[5+label]*objectness; if(confidence>=0.25f){ constfloatbx=ptr[0]; constfloatby=ptr[1]; constfloatbw=ptr[2]; constfloatbh=ptr[3]; Objectobj; //这里要减掉偏移值 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; }//iloop 对解析出的目标做非极大值抑制(NMS)操作后,检测结果如下图所示: 4a943c8a-58d6-11ed-a3b6-dac502259ad0.jpg

总结

本文以YOLOv5为例通过大量的代码一步步讲解如何使用TensorRT框架部署ONNX模型,主要目的是希望读者能够通过本文学习到TensorRT模型部署的基本流程,比如如何准备输入数据、如何调用API用模型做推理、如何解析模型的输出结果。如何部署YOLOv5模型并不是本文的重点,重点是要掌握使用TensorRT部署ONNX模型的基本方法,这样才会有举一反三的效果。

审核编辑:彭静
声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉
  • C++
    C++
    +关注

    关注

    21

    文章

    2094

    浏览量

    73442
  • 模型
    +关注

    关注

    1

    文章

    3097

    浏览量

    48624
  • 代码
    +关注

    关注

    30

    文章

    4708

    浏览量

    68174

原文标题:手把手教学,YOLOV5算法TensorRT部署流程

文章出处:【微信号:vision263com,微信公众号:新机器视觉】欢迎添加关注!文章转载请注明出处。

收藏 人收藏

    评论

    相关推荐

    深度探索ONNX模型部署 精选资料分享

    这篇文章从多个角度探索了ONNX,从ONNX的导出到ONNX和Caffe的对比,以及使用ONNX遭遇的困难以及一些解决办法,另...
    发表于 07-20 07:41

    ONNX的相关资料分享

    ONNX文件并生成特定平台和运行框架所支持的神经网络模型ONNX本身不是AI神经网络运行框架,只是AI神经网络
    发表于 11-05 06:45

    如何使用Paddle2ONNX模型转换工具将飞桨模型转换为ONNX模型

    如何使用Paddle2ONNX模型转换工具将飞桨模型转换为ONNX模型
    发表于 12-29 07:42

    EIQ onnx模型转换为tf-lite失败怎么解决?

    我们正在尝试将 tflite 框架与 npu 一起使用来进行机器学习。这是我们的步骤:1)用pytorch训练一个模型2) 以onnx格式导出模型3) eiq工具的covnert
    发表于 03-31 08:03

    Facebook推出ONNX,旨在为不同编程框架的神经网络创建共享模型

    今年9月,Facebook宣布推出“开源神经网络交换”(ONNX),呼吁其他公司加入,旨在为不同编程框架的神经网络创建共享模型。今天,Facebook联合AWS和微软宣布,在合作伙伴的支持下,第一个正式版本的
    的头像 发表于 12-28 16:12 4208次阅读

    嵌入式Linux平台部署AI神经网络模型Inference的方案

    ONNX文件并生成特定平台和运行框架所支持的神经网络模型ONNX本身不是AI神经网络运行框架,只是AI神经网络
    发表于 11-02 10:21 14次下载
    嵌入式Linux平台<b class='flag-5'>部署</b>AI神经网络<b class='flag-5'>模型</b>Inference的方案

    基于TensorRT完成NanoDet模型部署

    【GiantPandaCV导语】本文为大家介绍了一个TensorRT int8 量化部署 NanoDet 模型的教程,并开源了全部代码。主要是教你如何搭建tensorrt...
    发表于 01-25 19:04 0次下载
    基于<b class='flag-5'>TensorRT</b>完成NanoDet<b class='flag-5'>模型</b><b class='flag-5'>部署</b>

    使用Bottlerocket和Amazon EC2部署AI模型

      对于 AWS 上的人工智能推理部署,您可以利用 NVIDIA Triton 推理服务器 。使用开源推理服务软件在任何 CPU 或 CPU 基础设施上部署来自多个框架的经过培训的人工智能模型
    的头像 发表于 04-08 15:37 1047次阅读

    基于NVIDIA Triton的AI模型高效部署实践

    团队可以将各类框架(TensorFlowPyTorch、TensorRTONNX Runtime、MXNet、XGBoost 等或自定义框架后端)训练的 AI
    的头像 发表于 06-28 15:49 1591次阅读

    Pytorch转化ONNX过程代码实操

    一般来说转ONNX只是一个手段,在之后得到ONNX模型后还需要再将它做转换,比如转换到TensorRT上完成部署,或者有的人多加一步,从
    发表于 01-02 07:32 639次阅读

    ONNX格式模型部署兼容性框架介绍

      ONNXRUNTIME介绍 ONNX格式模型部署兼容性最强的框架 ONNXRUNTIME,基本上不会有算子不支持跟不兼容的情况出现,只要能导出O
    的头像 发表于 06-19 11:50 2350次阅读
    <b class='flag-5'>ONNX</b>格式<b class='flag-5'>模型</b><b class='flag-5'>部署</b>兼容性<b class='flag-5'>框架</b>介绍

    TorchVision框架模型导出并部署到ONNXRUNTIME C++全流程解析

    ONNXRUNTIME是主流的深度学习部署框架之一,支持ONNX格式模型在CPU、GPU、ARM等不同硬件平台上加速推理,支持C++、Python、Java、C#、JS等不同语言SDK
    的头像 发表于 07-13 14:46 1255次阅读
    TorchVision<b class='flag-5'>框架</b>下<b class='flag-5'>模型</b>导出并<b class='flag-5'>部署</b>到ONNXRUNTIME C++全流程解析

    三种主流模型部署框架YOLOv8推理演示

    深度学习模型部署有OpenVINO、ONNXRUNTIME、TensorRT三个主流框架,均支持Python与C++的SDK使用。对YOLOv5~YOLOv8的系列
    的头像 发表于 08-06 11:39 2560次阅读

    用STM32Cube.AI部署ONNX模型实操示例:风扇堵塞检测

    用STM32Cube.AI 部署ONNX 模型实操示例:风扇堵塞检测
    的头像 发表于 09-28 16:25 2341次阅读
    用STM32Cube.AI<b class='flag-5'>部署</b><b class='flag-5'>ONNX</b><b class='flag-5'>模型</b>实操示例:风扇堵塞检测

    基于Pytorch训练并部署ONNX模型在TDA4应用笔记

    电子发烧友网站提供《基于Pytorch训练并部署ONNX模型在TDA4应用笔记.pdf》资料免费下载
    发表于 09-11 09:24 0次下载
    基于Pytorch训练并<b class='flag-5'>部署</b><b class='flag-5'>ONNX</b><b class='flag-5'>模型</b>在TDA4应用笔记