TensorRT安裝

首先需要確保正確安裝CUDA，安裝后通過nvcc -V驗證是否安裝。

下載TensorRT
網址：https://developer.nvidia.com/nvidia-tensorrt-8x-download，下載最新版解壓

tar -xzvf TensorRT-8.4.0.6.Linux.x86_64-gnu.cuda-11.6.cudnn8.3.tar.gz

為了節省根目錄的內存，TensorRT我放在home下，添加環境變量

sudo vim ~/.bashrc export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:~/3D/TensorRT-8.4.0.6/lib source ~/.bashrc

接著，我們通過樣例程序判斷TRT能否正常工作。我們先是將sampleMNIST源碼進行編譯，此時會在bin目錄生成可執行文件，我們切換過去直接執行。

cd ~/3D/TensorRT-8.4.0.6/samples/sampleMNIST make cd ../../bin/ ./sample_mnist

如果輸出如下內容，最后顯示PASSED，表明樣例運行通過。

python支持

在上一部分中，雖然我們已經安裝了TensorRT，但是我們的Python環境還不能通過import tensorrt導入，所以需要通過安裝對應的.whl來實現。

pip install ../TensorRT-8.4.0.6/python/tensorrt-8.4.0.6-cp37-none-linux_x86_64.whl

ONNX部署

TensorRT是nvidia公司針對nvidia顯卡訓練的模型在特定平臺進行加速的推理優化庫，是一個c++庫，僅支持推理，不支持訓練；

進行推理，需要先創建IExecutionContext對象，要創建這個對象，就需要先創建一個ICudaEngine的對象（engine），兩種創建engine的方式：

• 使用模型文件創建engine，并可把創建的engine序列化后存儲到硬盤以便后面直接使用；
• 使用之前已經序列化存儲的engine，這種方式比較高效些，因為解析模型并生成engine還是挺慢的。

C++

TensorRT版本改動挺多的，大家可以直接查看API文檔

#include "NvInfer.h" #include "NvOnnxParser.h" #include "NvInferRuntimeCommon.h" #include <string> #include <iostream> #include <fstream>// Logger for TensorRT info/warning/errors class Logger : public nvinfer1::ILogger { public:Logger(Severity severity = Severity::kWARNING) : reportableSeverity(severity){}void log(Severity severity, char const* msg) noexcept// void log(Severity severity, const char* msg) noexcept{// suppress messages with severity enum value greater than the reportableif (severity > reportableSeverity)return;switch (severity){case Severity::kINTERNAL_ERROR:std::cerr << "INTERNAL_ERROR: ";break;case Severity::kERROR:std::cerr << "ERROR: ";break;case Severity::kWARNING:std::cerr << "WARNING: ";break;case Severity::kINFO:std::cerr << "INFO: ";break;default:std::cerr << "UNKNOWN: ";break;}std::cerr << msg << std::endl;}Severity reportableSeverity; };static Logger g_logger_;void onnxToTRTModel(const std::string &model_file, // name of the onnx modelnvinfer1::IHostMemory *&trt_model_stream) // output buffer for the TensorRT model {int verbosity = (int)nvinfer1::ILogger::Severity::kWARNING;// create the buildernvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(g_logger_);// 創建INetworkDefinition 對象nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U <<static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));// 創建解析器auto parser = nvonnxparser::createParser(*network, g_logger_);// 解析onnx文件，并填充網絡if (!parser->parseFromFile(model_file.c_str(), verbosity)){std::string msg("failed to parse onnx file");g_logger_.log(nvinfer1::ILogger::Severity::kERROR, msg.c_str());exit(EXIT_FAILURE);}// Build the enginebuilder->setMaxBatchSize(1);// 創建iBuilderConfig對象nvinfer1::IBuilderConfig* iBuilderConfig = builder->createBuilderConfig();// 設置engine可使用的最大GPU臨時值iBuilderConfig ->setMaxWorkspaceSize(1 << 20);nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network,*iBuilderConfig);// 將engine序列化，保存到文件中trt_model_stream = engine->serialize();// save enginestd::ofstream p("../model.trt", std::ios::binary);if (!p) {std::cerr << "could not open plan output file" << std::endl;}p.write(reinterpret_cast<const char *>(trt_model_stream->data()), trt_model_stream->size());parser->destroy();engine->destroy();network->destroy();builder->destroy();iBuilderConfig->destroy(); }int main() {nvinfer1::IHostMemory *trt_model_stream;onnxToTRTModel("../../config/pfe.onnx", trt_model_stream); }

CMakeLists.txt

cmake_minimum_required(VERSION 3.10) project(onnxToTensorRT) find_package(CUDA)SET(TENSORRT_PATH /home/xiaohu/3D/TensorRT-8.4.0.6/) SET(TENSORRT_LIB_PATH /home/xiaohu/3D/TensorRT-8.4.0.6/lib)include_directories(include${TENSORRT_PATH}/include${CUDA_INCLUDE_DIRS} )set(SOURCE_FILESonnxToTensorRT.cpp )file(GLOB TENSORRT_LIBS "${TENSORRT_LIB_PATH}/*.so") cuda_add_executable(${PROJECT_NAME} ${SOURCE_FILES}) target_link_libraries(${PROJECT_NAME}${TENSORRT_LIBS}${CUDA_LIBRARIES} )

python

############################################# # python api轉 ############################################# import sys import os import argparse import tensorrt as trtEXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)if __name__ == "__main__":parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference")parser.add_argument("--onnx_path", type=str,default='static_sim.onnx')parser.add_argument("--trt_path", type=str,default='static_sim.trt')args = parser.parse_args()onnx_file_path = args.onnx_pathengine_file_path = args.trt_pathprint('get start')TRT_LOGGER = trt.Logger()with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:config = builder.create_builder_config()config.max_workspace_size =( 1 << 30 ) * 2 # 2 GBbuilder.max_batch_size = 16config.set_flag(trt.BuilderFlag.FP16)# builder.fp16_mode = True# Parse model fileprint('Loading ONNX file from path {}...'.format(onnx_file_path))with open(onnx_file_path, 'rb') as model:print('Beginning ONNX file parsing')if not parser.parse(model.read()):print ('ERROR: Failed to parse the ONNX file.')for error in range(parser.num_errors):print (parser.get_error(error))print(f"raw shape of {network.get_input(0).name} is: ", network.get_input(0).shape)print(f"raw shape of {network.get_input(1).name} is: ", network.get_input(1).shape)print(f"raw shape of {network.get_input(2).name} is: ", network.get_input(2).shape)print('Completed parsing of ONNX file')print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))engine = builder.build_engine(network,config)print("Completed creating Engine")with open(engine_file_path, "wb") as f:f.write(engine.serialize())

trtexec

TensorRT-8.4.3.1/targets/x86_64-linux-gnu/bin/trtexec --onnx=static_sim.onnx --explicitBatch --saveEngine= static_sim.trt --workspace=1024 歡迎大家關注筆者，你的關注是我持續更博的最大動力

總結

以上是生活随笔為你收集整理的TensorRT8——ONNX转trt的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。