diff --git a/CMakeLists.txt b/CMakeLists.txt index 77425d3..f478cab 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.5) project (tkDNN) set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${CMAKE_CURRENT_SOURCE_DIR}/cmake) if(UNIX) -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -fPIC -Wno-deprecated-declarations -Wno-unused-variable") +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -fPIC -Wno-deprecated-declarations -Wno-unused-variable") endif() if(WIN32) set(CMAKE_CXX_STANDARD 14) diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h index 4c6c816..ca012fb 100644 --- a/include/tkDNN/NetworkRT.h +++ b/include/tkDNN/NetworkRT.h @@ -6,6 +6,7 @@ #include "Network.h" #include "Layer.h" #include "NvInfer.h" +#include namespace tk { namespace dnn { @@ -59,7 +60,8 @@ public: #if NV_TENSORRT_MAJOR >= 6 nvinfer1::IBuilderConfig *configRT; #endif - nvinfer1::ICudaEngine *engineRT; + std::shared_ptr engineRT; + //nvinfer1::ICudaEngine *engineRT; nvinfer1::IExecutionContext *contextRT; const static int MAX_BUFFERS_RT = 10; @@ -114,6 +116,9 @@ public: bool serialize(const char *filename); bool deserialize(const char *filename); + + + }; }} diff --git a/include/tkDNN/utils.h b/include/tkDNN/utils.h index 1404509..edc770e 100644 --- a/include/tkDNN/utils.h +++ b/include/tkDNN/utils.h @@ -110,6 +110,17 @@ double t_ns = time_ms.count(); FatalError(_error.str()); \ } \ } +struct InferDeleter +{ + template + void operator()(T* obj) const + { + if (obj) + { + obj->destroy(); + } + } +}; typedef enum { ERROR_CUDNN = 2, diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 501ade4..4006caf 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -137,9 +137,11 @@ NetworkRT::NetworkRT(Network *net, const char *name) { printCudaMemUsage(); std::cout<<"Building tensorRT cuda engine...\n"; #if NV_TENSORRT_MAJOR >= 6 - engineRT = builderRT->buildEngineWithConfig(*networkRT, *configRT); + //engineRT = builderRT->buildEngineWithConfig(*networkRT, *configRT); + engineRT = std::shared_ptr(builderRT->buildEngineWithConfig(*networkRT,*configRT),InferDeleter()); #else - engineRT = builderRT->buildCudaEngine(*networkRT); + //engineRT = builderRT->buildCudaEngine(*networkRT); + engineRT = std::shared_ptr(builderRT->buildCudaEngine(*networkRT)); #endif if(engineRT == nullptr) FatalError("cloud not build cuda engine") @@ -561,7 +563,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, DeformConv2d *l) { IPluginLayer *lRT = networkRT->addPlugin(inputs, 2, *plugin); checkNULL(lRT); lRT->setName( ("Deformable" + std::to_string(l->id)).c_str() ); - delete(inputs); + delete[](inputs); // batchnorm void *bias_b, *power_b, *mean_b, *variance_b, *scales_b; if(dtRT == DataType::kHALF) { @@ -629,7 +631,8 @@ bool NetworkRT::deserialize(const char *filename) { pluginFactory = new PluginFactory(); runtimeRT = createInferRuntime(loggerRT); - engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size, (IPluginFactory *) pluginFactory); + //engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size, (IPluginFactory *) pluginFactory); + engineRT = std::shared_ptr(runtimeRT->deserializeCudaEngine(gieModelStream,size,(IPluginFactory*)pluginFactory),InferDeleter()); //if (gieModelStream) delete [] gieModelStream; return true;