From 30de4562bcc9abd3024da53ccd0c145e29d0c1c7 Mon Sep 17 00:00:00 2001 From: tsieyy <77835822+tsieyy@users.noreply.github.com> Date: Wed, 19 Apr 2023 12:45:43 +0800 Subject: [PATCH] Adapt ufld model to tensorrt8 (#1288) * make ufld adapt tensorrt8 * make ufld adapt tensorrt8 and tensorrt7 --- ufld/CMakeLists.txt | 9 +++++++-- ufld/gen_wts.py | 2 +- ufld/lane_det.cpp | 19 ++++++++++--------- ufld/logging.h | 3 ++- ufld/macros.h | 27 +++++++++++++++++++++++++++ 5 files changed, 47 insertions(+), 13 deletions(-) create mode 100644 ufld/macros.h diff --git a/ufld/CMakeLists.txt b/ufld/CMakeLists.txt index 3454e37..68f927d 100644 --- a/ufld/CMakeLists.txt +++ b/ufld/CMakeLists.txt @@ -8,9 +8,14 @@ option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) set(CMAKE_CXX_STANDARD 11) set(CMAKE_BUILD_TYPE Debug) -find_package(CUDA REQUIRED) +# cuda directory +include_directories(/usr/local/cuda/include/) +link_directories(/usr/local/cuda/lib64/) + +# tensorrt +#include_directories(/workspace/TensorRT-7.2.3.4/include/) +#link_directories(/workspace/TensorRT-7.2.3.4/lib/) -include_directories(${PROJECT_SOURCE_DIR}/include) find_package(OpenCV) include_directories(${OpenCV_INCLUDE_DIRS}) diff --git a/ufld/gen_wts.py b/ufld/gen_wts.py index 33abde6..dd13b75 100644 --- a/ufld/gen_wts.py +++ b/ufld/gen_wts.py @@ -4,7 +4,7 @@ import struct from model.model import parsingNet # Initialize -model = parsingNet(pretrained = False, backbone='18', cls_dim = (101, 56, 4), use_aux=False).cuda() +model = parsingNet(pretrained = False, backbone='18', cls_dim = (101, 56, 4), use_aux=False) device = 'cpu' # Load model state_dict = torch.load('tusimple_18.pth', map_location='cpu')['model'] diff --git a/ufld/lane_det.cpp b/ufld/lane_det.cpp index 5ab376e..aab1af4 100644 --- a/ufld/lane_det.cpp +++ b/ufld/lane_det.cpp @@ -21,8 +21,8 @@ const char* OUTPUT_BLOB_NAME = "prob"; static Logger gLogger; // Creat the engine using only the API and not any parser. -ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType dt) { - INetworkDefinition* network = builder->createNetwork(); +ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder,IBuilderConfig* builderConfig, DataType dt) { + INetworkDefinition* network = builder->createNetworkV2(0U); Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{INPUT_C, INPUT_H, INPUT_W }); @@ -84,7 +84,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataTyp permute0->setReshapeDimensions( Dims2{1, 1800}); auto fcwts0 = network->addConstant(nvinfer1::Dims2(2048, 1800), weightMap["cls.0.weight"]); - auto matrixMultLayer0 = network->addMatrixMultiply(*permute0->getOutput(0), false, *fcwts0->getOutput(0), true); + auto matrixMultLayer0 = network->addMatrixMultiply(*permute0->getOutput(0), MatrixOperation::kNONE, *fcwts0->getOutput(0), MatrixOperation::kTRANSPOSE); assert(matrixMultLayer0 != nullptr); // Add elementwise layer for adding bias @@ -96,7 +96,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataTyp auto relu = network->addActivation(*addBiasLayer0->getOutput(0), ActivationType::kRELU); auto fcwts1 = network->addConstant(nvinfer1::Dims2(22624, 2048), weightMap["cls.2.weight"]); - auto matrixMultLayer1 = network->addMatrixMultiply(*relu->getOutput(0), false, *fcwts1->getOutput(0), true); + auto matrixMultLayer1 = network->addMatrixMultiply(*relu->getOutput(0), MatrixOperation::kNONE, *fcwts1->getOutput(0), MatrixOperation::kTRANSPOSE); assert(matrixMultLayer1 != nullptr); // Add elementwise layer for adding bias @@ -114,17 +114,18 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataTyp // Build engine builder->setMaxBatchSize(maxBatchSize); - builder->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB + builderConfig->setMaxWorkspaceSize(16 * (1 << 20));// 16MB + #ifdef USE_FP16 if(builder->platformHasFastFp16()) { std::cout << "Platform supports fp16 mode and use it !!!" << std::endl; - builder->setFp16Mode(true); + builderConfig->setFlag(BuilderFlag::kFP16); } else { std::cout << "Platform doesn't support fp16 mode so you can't use it !!!" << std::endl; } #endif std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildCudaEngine(*network); + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *builderConfig); std::cout << "Build engine successfully!" << std::endl; // Don't need the network any more @@ -142,9 +143,9 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataTyp void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) { // Create builder IBuilder* builder = createInferBuilder(gLogger); - + IBuilderConfig* builderConfig = builder->createBuilderConfig(); // Create model to populate the network, then set the outputs and create an engine - ICudaEngine* engine = createEngine(maxBatchSize, builder, DataType::kFLOAT); + ICudaEngine* engine = createEngine(maxBatchSize, builder, builderConfig, DataType::kFLOAT); assert(engine != nullptr); // Serialize the engine diff --git a/ufld/logging.h b/ufld/logging.h index b817098..4f4a83e 100644 --- a/ufld/logging.h +++ b/ufld/logging.h @@ -9,6 +9,7 @@ #include #include #include +#include "macros.h" using Severity = nvinfer1::ILogger::Severity; @@ -220,7 +221,7 @@ public: //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the //! inheritance from nvinfer1::ILogger //! - void log(Severity severity, const char* msg) override + void log(Severity severity, const char* msg) TRT_NOEXCEPT override { LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; } diff --git a/ufld/macros.h b/ufld/macros.h new file mode 100644 index 0000000..0e5b71b --- /dev/null +++ b/ufld/macros.h @@ -0,0 +1,27 @@ +#ifndef __MACROS_H +#define __MACROS_H + +#ifdef API_EXPORTS +#if defined(_MSC_VER) +#define API __declspec(dllexport) +#else +#define API __attribute__((visibility("default"))) +#endif +#else + +#if defined(_MSC_VER) +#define API __declspec(dllimport) +#else +#define API +#endif +#endif // API_EXPORTS + +#if NV_TENSORRT_MAJOR >= 8 +#define TRT_NOEXCEPT noexcept +#define TRT_CONST_ENQUEUE const +#else +#define TRT_NOEXCEPT +#define TRT_CONST_ENQUEUE +#endif + +#endif // __MACROS_H