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

導讀 模型部署作為算法模型落地的最后一步，在人工智能產業化過程中是非常關鍵的步驟，而目標檢測作為計算機視覺三大基礎任務之一，眾多的業務功能都要在檢測的基礎之上完成，本文提供了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(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?< MyLoggerlogger; nvinfer1::IBuilder*builder=nvinfer1::createInferBuilder(logger); 創建IBuilder對象后，優化一個模型的第一步是要構建模型的網絡結構。

constuint32_texplicit_batch=1U<( 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;igetNbErrors();++i){ std::cout<getError(i)->desc()< nvinfer1::IBuilderConfig*config=builder->createBuilderConfig(); config->setMemoryPoolLimit(nvinfer1::kWORKSPACE,1U<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< 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<(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框架進行模型推理的完整流程：

對輸入圖像數據做與模型訓練時一樣的預處理操作。

把模型的輸入數據從CPU拷貝到GPU中。

調用模型推理接口進行推理。

把模型的輸出數據從GPU拷貝到CPU中。

對模型的輸出結果進行解析，進行必要的后處理后得到最終的結果。

由于模型的推理是在GPU上進行的，所以會存在搬運輸入、輸出數據的操作，因此有必要在GPU上創建內存區域用于存放輸入、輸出數據。模型輸入、輸出的尺寸可以通過ICudaEngine對象的接口來獲取，根據這些信息我們可以先為模型分配輸入、輸出緩存區。

void*buffers[2]; //獲取模型輸入尺寸并分配GPU內存 nvinfer1::Dimsinput_dim=engine->getBindingDimensions(0); intinput_size=1; for(intj=0;jgetBindingDimensions(1); intoutput_size=1; for(intj=0;j 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模型的基本流程。 如上圖所示，主要步驟如下：

實例化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(inti=0;i=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）操作后，檢測結果如下圖所示：

總結

本文以YOLOv5為例通過大量的代碼一步步講解如何使用TensorRT框架部署ONNX模型，主要目的是希望讀者能夠通過本文學習到TensorRT模型部署的基本流程，比如如何準備輸入數據、如何調用API用模型做推理、如何解析模型的輸出結果。如何部署YOLOv5模型并不是本文的重點，重點是要掌握使用TensorRT部署ONNX模型的基本方法，這樣才會有舉一反三的效果。

審核編輯：彭靜