TensorRT安裝

首先需要確保正確安裝CUDA，安裝后通過(guò)nvcc -V驗(yàn)證是否安裝。

下載TensorRT
網(wǎng)址：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

為了節(jié)省根目錄的內(nèi)存，TensorRT我放在home下，添加環(huán)境變量

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

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

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

如果輸出如下內(nèi)容，最后顯示PASSED，表明樣例運(yùn)行通過(guò)。

python支持

在上一部分中，雖然我們已經(jīng)安裝了TensorRT，但是我們的Python環(huán)境還不能通過(guò)import tensorrt導(dǎo)入，所以需要通過(guò)安裝對(duì)應(yīng)的.whl來(lái)實(shí)現(xiàn)。

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

ONNX部署

TensorRT是nvidia公司針對(duì)nvidia顯卡訓(xùn)練的模型在特定平臺(tái)進(jìn)行加速的推理優(yōu)化庫(kù)，是一個(gè)c++庫(kù)，僅支持推理，不支持訓(xùn)練；

進(jìn)行推理，需要先創(chuàng)建IExecutionContext對(duì)象，要?jiǎng)?chuàng)建這個(gè)對(duì)象，就需要先創(chuàng)建一個(gè)ICudaEngine的對(duì)象（engine），兩種創(chuàng)建engine的方式：

• 使用模型文件創(chuàng)建engine，并可把創(chuàng)建的engine序列化后存儲(chǔ)到硬盤以便后面直接使用；
• 使用之前已經(jīng)序列化存儲(chǔ)的engine，這種方式比較高效些，因?yàn)榻馕瞿Ｐ筒⑸蒭ngine還是挺慢的。

C++

TensorRT版本改動(dòng)挺多的，大家可以直接查看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_);// 創(chuàng)建INetworkDefinition 對(duì)象nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U <<static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));// 創(chuàng)建解析器auto parser = nvonnxparser::createParser(*network, g_logger_);// 解析onnx文件，并填充網(wǎng)絡(luò)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);// 創(chuàng)建iBuilderConfig對(duì)象nvinfer1::IBuilderConfig* iBuilderConfig = builder->createBuilderConfig();// 設(shè)置engine可使用的最大GPU臨時(shí)值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轉(zhuǎn) ############################################# 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 歡迎大家關(guān)注筆者，你的關(guān)注是我持續(xù)更博的最大動(dòng)力

總結(jié)

以上是生活随笔為你收集整理的TensorRT8——ONNX转trt的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。