From 4d8f99b44143c8c140e98a2efc4b04e198f49d85 Mon Sep 17 00:00:00 2001 From: perseusdg Date: Sun, 9 Jan 2022 19:24:38 +0530 Subject: [PATCH] depth->tensorrt8 patches --- CMakeLists.txt | 5 +- include/tkDNN/NetworkRT.h | 2 + include/tkDNN/pluginsRT/ReflectionPadding.h | 2 + src/NetworkRT.cpp | 55 +++++++++++++++------ 4 files changed, 46 insertions(+), 18 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 7cc9e33..d919b46 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -85,12 +85,14 @@ endif() find_package(CUDNN REQUIRED) include_directories(${CUDNN_INCLUDE_DIR}) +find_package(yaml-cpp REQUIRED) + # compile file(GLOB tkdnn_CUSRC "src/kernels/*.cu" "src/sorting.cu" "src/pluginsRT/*.cpp") cuda_include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${CUDNN_INCLUDE_DIRS}) cuda_add_library(kernels SHARED ${tkdnn_CUSRC}) -target_link_libraries(kernels ${CUDA_CUBLAS_LIBRARIES} ${CUDA_LIBRARIES} ${CUDNN_LIBRARIES}) +target_link_libraries(kernels ${CUDA_CUBLAS_LIBRARIES} ${CUDA_LIBRARIES} ${CUDNN_LIBRARIES} yaml-cpp) @@ -120,7 +122,6 @@ endif() # endif() # gives problems in cross-compiling, probably malformed cmake config -find_package(yaml-cpp REQUIRED) #------------------------------------------------------------------------------- # Build Libraries diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h index 571d127..b859074 100644 --- a/include/tkDNN/NetworkRT.h +++ b/include/tkDNN/NetworkRT.h @@ -23,6 +23,8 @@ #include #include #include +#include +#include diff --git a/include/tkDNN/pluginsRT/ReflectionPadding.h b/include/tkDNN/pluginsRT/ReflectionPadding.h index 894ed98..7b13710 100644 --- a/include/tkDNN/pluginsRT/ReflectionPadding.h +++ b/include/tkDNN/pluginsRT/ReflectionPadding.h @@ -94,6 +94,8 @@ namespace nvinfer1{ static std::vector mPluginAttributes; std::string mPluginNamespace; }; + + REGISTER_TENSORRT_PLUGIN(ReflectionPaddingRTPluginCreator); }; #endif diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 6192006..3f086bd 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -474,23 +474,46 @@ ILayer* NetworkRT::convert_layer(ITensor *input,Padding *l){ #else //todo add PADDING_MODE_CONSTANT AND PADDING_MODE_ZERO for tensorrt versions < 8.2 if(l->padding_mode == PADDING_MODE_REFLECTION){ - auto creator = getPluginRegistry()->getPluginCreator("ReflectionPaddingRT_tkDNN","1"); - std::vector mPluginAttributes; - PluginFieldCollection mFC{}; - mPluginAttributes.emplace_back(PluginField("padH",&l->paddingH,PluginFieldType::kINT32,1)); - mPluginAttributes.emplace_back(PluginField("padW",&l->paddingW,PluginFieldType::kINT32,1)); - mPluginAttributes.emplace_back(PluginField("inputH",&l->input_dim.h,PluginFieldType::kINT32,1)); - mPluginAttributes.emplace_back(PluginField("inputW",&l->input_dim.w,PluginFieldType::kINT32,1)); - mPluginAttributes.emplace_back(PluginField("outputH",&l->output_dim.h,PluginFieldType::kINT32,1)); - mPluginAttributes.emplace_back(PluginField("outputW",&l->output_dim.w,PluginFieldType::kINT32,1)); - mPluginAttributes.emplace_back(PluginField("n",&l->input_dim.n,PluginFieldType::kINT32,1)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); - auto *plugin = creator->createPlugin(l->getLayerName().c_str(),&mFC); + auto creator = getPluginRegistry()->getPluginCreator("ReflectionPaddingRT_tkDNN","1"); + std::vector mPluginAttributes; + PluginFieldCollection mFC{}; + mPluginAttributes.emplace_back(PluginField("padH",&l->paddingH,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("padW",&l->paddingW,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("inputH",&l->input_dim.h,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("inputW",&l->input_dim.w,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("outputH",&l->output_dim.h,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("outputW",&l->output_dim.w,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("n",&l->input_dim.n,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("c",&l->input_dim.c,PluginFieldType::kINT32,1)); + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); + auto *plugin = creator->createPlugin(l->getLayerName().c_str(),&mFC); + auto *lRT = networkRT->addPluginV2(&input, 1, *plugin); + checkNULL(lRT); + return lRT; + }else if(l->padding_mode == PADDING_MODE_CONSTANT || l->padding_mode == PADDING_MODE_ZERO){ + auto creator = getPluginRegistry()->getPluginCreator("ConstantPaddingRT_tkDNN","1"); + std::vector mPluginAttributes; + PluginFieldCollection mFC{}; + mPluginAttributes.emplace_back(PluginField("padH",&l->paddingH,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("padW",&l->paddingW,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("inputH",&l->input_dim.h,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("inputW",&l->input_dim.w,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("outputH",&l->output_dim.h,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("outputW",&l->output_dim.w,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("n",&l->input_dim.n,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("c",&l->input_dim.c,PluginFieldType::kINT32,1)); + mPluginAttributes.emplace_back(PluginField("constant",&l->constant,PluginFieldType::kFLOAT32,1)); + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); + auto *plugin = creator->createPlugin(l->getLayerName().c_str(),&mFC); + auto *lRT = networkRT->addPluginV2(&input,1,*plugin); + checkNULL(lRT); + return lRT; } - auto *lRT = networkRT->addPluginV2(&input, 1, *plugin); - checkNULL(lRT); - return lRT; + + return nullptr; + #endif }