From 1ad7f9584b29b8a072ef3f49e3d7ba15814b9b4d Mon Sep 17 00:00:00 2001 From: l <1430244438@qq.com> Date: Mon, 23 Nov 2020 22:11:09 +0800 Subject: [PATCH] add surport for tensorrt7 --- cmake/tensorrt-config.cmake | 8 +++--- include/netTensorRT.hpp | 2 +- src/netTensorRT.cpp | 51 +++++++++++++++++++++++++++++-------- 3 files changed, 45 insertions(+), 16 deletions(-) diff --git a/cmake/tensorrt-config.cmake b/cmake/tensorrt-config.cmake index 014deec..2f78b45 100644 --- a/cmake/tensorrt-config.cmake +++ b/cmake/tensorrt-config.cmake @@ -6,10 +6,10 @@ find_library(NVINFER NAMES nvinfer) find_library(NVINFERPLUGIN NAMES nvinfer_plugin) find_library(NVPARSERS NAMES nvparsers) find_library(NVONNXPARSER NAMES nvonnxparser) -find_library(NVONNXPARSERRUNTIME NAMES nvonnxparser_runtime) +# find_library(NVONNXPARSERRUNTIME NAMES nvonnxparser_runtime) # If it is ALL there, export libraries as a single package -if(CUDA_FOUND AND NVINFER AND NVINFERPLUGIN AND NVPARSERS AND NVONNXPARSER AND NVONNXPARSERRUNTIME) +if(CUDA_FOUND AND NVINFER AND NVINFERPLUGIN AND NVPARSERS AND NVONNXPARSER) message("TensorRT available!") message("CUDA Libs: ${CUDA_LIBRARIES}") message("CUDA Headers: ${CUDA_INCLUDE_DIRS}") @@ -17,8 +17,8 @@ if(CUDA_FOUND AND NVINFER AND NVINFERPLUGIN AND NVPARSERS AND NVONNXPARSER AND N message("NVINFERPLUGIN: ${NVINFERPLUGIN}") message("NVPARSERS: ${NVPARSERS}") message("NVONNXPARSER: ${NVONNXPARSER}") - message("NVONNXPARSERRUNTIME: ${NVONNXPARSERRUNTIME}") - list(APPEND TENSORRT_LIBRARIES ${CUDA_LIBRARIES} nvinfer nvinfer_plugin nvparsers nvonnxparser nvonnxparser_runtime) + # message("NVONNXPARSERRUNTIME: ${NVONNXPARSERRUNTIME}") + list(APPEND TENSORRT_LIBRARIES ${CUDA_LIBRARIES} nvinfer nvinfer_plugin nvparsers nvonnxparser) message("All togheter now (libs): ${TENSORRT_LIBRARIES}") list(APPEND TENSORRT_INCLUDE_DIRS ${CUDA_INCLUDE_DIRS}) message("All togheter now (inc): ${TENSORRT_INCLUDE_DIRS}") diff --git a/include/netTensorRT.hpp b/include/netTensorRT.hpp index dce3259..d95cd71 100644 --- a/include/netTensorRT.hpp +++ b/include/netTensorRT.hpp @@ -8,7 +8,7 @@ // For plugin factory #include #include -#include +// #include #include #include #include diff --git a/src/netTensorRT.cpp b/src/netTensorRT.cpp index d889a43..7d23509 100644 --- a/src/netTensorRT.cpp +++ b/src/netTensorRT.cpp @@ -449,12 +449,12 @@ void NetTensorRT::deserializeEngine(const std::string& engine_path) { << std::endl; // because I use onnx-tensorRT i have to use their plugin factory - nvonnxparser::IPluginFactory* plug_fact = - nvonnxparser::createPluginFactory(_gLogger); + // nvonnxparser::IPluginFactory* plug_fact = + // nvonnxparser::createPluginFactory(_gLogger); // Now deserialize - _engine = infer->deserializeCudaEngine(modelMem, modelSize, plug_fact); - + _engine = infer->deserializeCudaEngine(modelMem, modelSize, nullptr); + infer->destroy(); free(modelMem); if (_engine) { std::cerr << "Created engine!" << std::endl; @@ -500,19 +500,41 @@ void NetTensorRT::generateEngine(const std::string& onnx_path) { // create inference builder IBuilder* builder = createInferBuilder(_gLogger); + if (builder) + { + std::cout << "create builder succeed!\n" + << std::endl; + } + else + { + throw std::runtime_error("ERROR: could not create builder."); + } + // BATCH SIZE IS ALWAYS ONE + builder->setMaxBatchSize(1); + IBuilderConfig *builderconfig = builder->createBuilderConfig(); + if (builderconfig) + { + std::cout << "create builderconfig succeed!\n" + << std::endl; + } + else + { + throw std::runtime_error("ERROR: could not create builderconfig."); + } // set optimization parameters here // CAN I DO HALF PRECISION (and report to user) std::cout << "Platform "; if (builder->platformHasFastFp16()) { std::cout << "HAS "; - builder->setFp16Mode(true); + // builder->setFp16Mode(true); + builderconfig->setFlag(BuilderFlag::kFP16); } else { std::cout << "DOESN'T HAVE "; - builder->setFp16Mode(false); + // builder->setFp16Mode(false); } std::cout << "fp16 support." << std::endl; // BATCH SIZE IS ALWAYS ONE - builder->setMaxBatchSize(1); + // builder->setMaxBatchSize(1); // if using DLA, set the desired core before deserialization occurs #if NV_TENSORRT_MAJOR >= 5 && \ @@ -534,7 +556,9 @@ void NetTensorRT::generateEngine(const std::string& onnx_path) { #endif // create a network builder - INetworkDefinition* network = builder->createNetwork(); + // INetworkDefinition* network = builder->createNetwork(); + const auto explicitBatch = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); + INetworkDefinition *network = builder->createNetworkV2(explicitBatch); // generate a parser to get weights from onnx file nvonnxparser::IParser* parser = @@ -553,10 +577,12 @@ void NetTensorRT::generateEngine(const std::string& onnx_path) { for (unsigned long ws_size = MAX_WORKSPACE_SIZE; ws_size >= MIN_WORKSPACE_SIZE; ws_size /= 2) { // set size - builder->setMaxWorkspaceSize(ws_size); + // builder->setMaxWorkspaceSize(ws_size); + builderconfig->setMaxWorkspaceSize(ws_size); // try to build - _engine = builder->buildCudaEngine(*network); + // _engine = builder->buildCudaEngine(*network); + _engine = builder->buildEngineWithConfig(*network, *builderconfig); if (!_engine) { std::cerr << "Failure creating engine from ONNX model" << std::endl << "Current trial size is " << ws_size << std::endl; @@ -567,7 +593,10 @@ void NetTensorRT::generateEngine(const std::string& onnx_path) { break; } } - + builder->destroy(); + builderconfig->destroy(); + parser->destroy(); + network->destroy(); // final check if (!_engine) { throw std::runtime_error("ERROR: could not create engine from ONNX.");