diff --git a/CMakeLists.txt b/CMakeLists.txt index c10f6ac..bfdb426 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -51,13 +51,13 @@ endif() if(WIN32) if(CMAKE_BUILD_TYPE MATCHES Release) - set(CMAKE_CXX_FLAGS "/O2 /FS /EHsc") - set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32 -arch=sm_61) + set(CMAKE_CXX_FLAGS "/O2 /FS /EHsc /Md") + set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32) endif() if(CMAKE_BUILD_TYPE MATCHES Debug) - set(CMAKE_CXX_FLAGS "/Od /FS /EHsc") - set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32 -arch=sm_61 -G -g) + set(CMAKE_CXX_FLAGS "/Od /FS /EHsc /MDd") + set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32 -G -g) endif() set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON) endif(WIN32) @@ -92,7 +92,7 @@ include_directories(${CUDNN_INCLUDE_DIR}) 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_LIBRAY} ${CUDA_LIBRARIES} ${CUDNN_LIBRARIES}) +target_link_libraries(kernels ${CUDA_CUBLAS_LIBRARIES} ${CUDA_LIBRARIES} ${CUDNN_LIBRARIES}) @@ -121,7 +121,7 @@ set(tkdnn_LIBS kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} ${CUDNN_LIBRAR set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}") include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${OPENCV_INCLUDE_DIRS} ${NVINFER_INCLUDES}) add_library(tkDNN SHARED ${tkdnn_SRC}) -target_link_libraries(tkDNN ${tkdnn_LIBS} ${CUDA_CUBLAS_LIBRAY}) +target_link_libraries(tkDNN ${tkdnn_LIBS} ${CUDA_CUBLAS_LIBRARIES}) #static #add_library(tkDNN_static STATIC ${tkdnn_SRC}) diff --git a/include/tkDNN/Layer.h b/include/tkDNN/Layer.h index 0c1c5ca..d08e517 100644 --- a/include/tkDNN/Layer.h +++ b/include/tkDNN/Layer.h @@ -423,6 +423,8 @@ public: virtual layerType_t getLayerType() { return LAYER_FLATTEN; }; virtual dnnType* infer(dataDim_t &dim, dnnType* srcData); + + int c, h, w, rows, cols; }; /** diff --git a/include/tkDNN/pluginsRT/FlattenConcatRT.h b/include/tkDNN/pluginsRT/FlattenConcatRT.h index b1e5327..aa8f078 100644 --- a/include/tkDNN/pluginsRT/FlattenConcatRT.h +++ b/include/tkDNN/pluginsRT/FlattenConcatRT.h @@ -3,10 +3,10 @@ #include #include namespace nvinfer1 { - class FlattenConcatRT : public IPluginV2 { + class FlattenConcatRT : public IPluginV2IOExt { public: - FlattenConcatRT() ; + FlattenConcatRT(int c,int h,int w,int rows,int cols) ; FlattenConcatRT(const void *data, size_t length) ; @@ -16,9 +16,6 @@ namespace nvinfer1 { Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ; - void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type, - PluginFormat format, int maxBatchSize) NOEXCEPT override ; - int initialize() NOEXCEPT override ; void terminate() NOEXCEPT override ; @@ -37,8 +34,6 @@ namespace nvinfer1 { void destroy() NOEXCEPT override ; - bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ; - const char *getPluginType() const NOEXCEPT override ; const char *getPluginVersion() const NOEXCEPT override; @@ -47,7 +42,21 @@ namespace nvinfer1 { void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ; - IPluginV2 *clone() const NOEXCEPT override ; + IPluginV2IOExt *clone() const NOEXCEPT override ; + + DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override; + + void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) NOEXCEPT override; + + void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override; + + bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override; + + bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override; + + bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const NOEXCEPT override; + + void detachFromContext() NOEXCEPT override; int c, h, w; int rows, cols; @@ -65,9 +74,9 @@ namespace nvinfer1 { const char *getPluginNamespace() const NOEXCEPT override ; - IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ; + IPluginV2IOExt *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ; - IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ; + IPluginV2IOExt *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ; const char *getPluginName() const NOEXCEPT override ; diff --git a/src/Flatten.cpp b/src/Flatten.cpp index 8906c88..0eb690c 100644 --- a/src/Flatten.cpp +++ b/src/Flatten.cpp @@ -15,6 +15,11 @@ Flatten::Flatten(Network *net) : Layer(net) { output_dim.w = 1; output_dim.l = 1; + this->h = 1; + this->w = 1; + this->rows = input_dim.w; + this->cols = input_dim.h * input_dim.c; + this->c = input_dim.w * input_dim.h * input_dim.c; } Flatten::~Flatten() { diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 4b3d532..0537e6a 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -18,9 +18,9 @@ using namespace nvinfer1; // Logger for info/warning/errors class Logger : public ILogger { void log(Severity severity, const char* msg) NOEXCEPT override { -#ifdef DEBUG +//#ifdef DEBUG std::cout <<"TENSORRT LOG: "<< msg << std::endl; -#endif +//#endif } } loggerRT; @@ -473,7 +473,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Route *l) { ILayer* NetworkRT::convert_layer(ITensor *input, Flatten *l) { - IPluginV2 *plugin = new FlattenConcatRT(); + IPluginV2IOExt *plugin = new FlattenConcatRT(l->c,l->h,l->w,l->rows,l->cols); IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin); checkNULL(lRT); return lRT; diff --git a/src/pluginsRT/FlattenConcatRT.cpp b/src/pluginsRT/FlattenConcatRT.cpp index 6687b58..6c54e47 100644 --- a/src/pluginsRT/FlattenConcatRT.cpp +++ b/src/pluginsRT/FlattenConcatRT.cpp @@ -4,12 +4,17 @@ using namespace nvinfer1; std::vector FlattenConcatRTPluginCreator::mPluginAttributes; PluginFieldCollection FlattenConcatRTPluginCreator::mFC{}; -FlattenConcatRT::FlattenConcatRT() { +FlattenConcatRT::FlattenConcatRT(int c, int h, int w, int rows, int cols) { stat = cublasCreate(&handle); if (stat != CUBLAS_STATUS_SUCCESS) { printf ("CUBLAS initialization failed\n"); return; } + this->c = c; + this->h = h; + this->w = w; + this->rows = rows; + this->cols = cols; } FlattenConcatRT::FlattenConcatRT(const void *data, size_t length) { @@ -32,16 +37,6 @@ Dims FlattenConcatRT::getOutputDimensions(int index, const Dims *inputs, int nbI return Dims3{ inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2], 1, 1}; } -void FlattenConcatRT::configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, - DataType type, PluginFormat format, int maxBatchSize) NOEXCEPT { - assert(nbOutputs == 1 && nbInputs ==1); - rows = inputDims[0].d[0]; - cols = inputDims[0].d[1] * inputDims[0].d[2]; - c = inputDims[0].d[0] * inputDims[0].d[1] * inputDims[0].d[2]; - h = 1; - w = 1; -} - int FlattenConcatRT::initialize() NOEXCEPT { return 0; } @@ -107,9 +102,7 @@ void FlattenConcatRT::destroy() NOEXCEPT { delete this; } -bool FlattenConcatRT::supportsFormat(DataType type, PluginFormat format) const NOEXCEPT { - return true; -} + const char *FlattenConcatRT::getPluginType() const NOEXCEPT { return "FlattenConcatRT_tkDNN"; @@ -127,12 +120,44 @@ void FlattenConcatRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT { mPluginNamespace = pluginNamespace; } -IPluginV2 *FlattenConcatRT::clone() const NOEXCEPT { - auto *p = new FlattenConcatRT(); +IPluginV2IOExt *FlattenConcatRT::clone() const NOEXCEPT { + auto* p = new FlattenConcatRT(c, h, w, rows, cols); p->setPluginNamespace(mPluginNamespace.c_str()); return p; } +DataType FlattenConcatRT::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT +{ + return DataType::kFLOAT; +} + +void FlattenConcatRT::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) NOEXCEPT +{ +} + +void FlattenConcatRT::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT +{ +} + +bool FlattenConcatRT::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT +{ + return false; +} + +bool FlattenConcatRT::canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT +{ + return false; +} + +bool FlattenConcatRT::supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const NOEXCEPT +{ + return true; +} + +void FlattenConcatRT::detachFromContext() NOEXCEPT +{ +} + FlattenConcatRTPluginCreator::FlattenConcatRTPluginCreator() { mPluginAttributes.clear(); mFC.nbFields = mPluginAttributes.size(); @@ -147,15 +172,21 @@ const char *FlattenConcatRTPluginCreator::getPluginNamespace() const NOEXCEPT { return mPluginNamespace.c_str(); } -IPluginV2 *FlattenConcatRTPluginCreator::deserializePlugin(const char *name, const void *serialData, +IPluginV2IOExt *FlattenConcatRTPluginCreator::deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT { auto *pluginObj = new FlattenConcatRT(serialData,serialLength); pluginObj->setPluginNamespace(mPluginNamespace.c_str()); return pluginObj; } -IPluginV2 *FlattenConcatRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT { - auto *pluginObj = new FlattenConcatRT(); +IPluginV2IOExt *FlattenConcatRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT { + const PluginField* fields = fc->fields; + int c = *(static_cast(fields[0].data)); + int h = *(static_cast(fields[1].data)); + int w = *(static_cast(fields[2].data)); + int rows = *(static_cast(fields[3].data)); + int cols = *(static_cast(fields[4].data)); + auto* pluginObj = new FlattenConcatRT(c, h, w, rows, cols); pluginObj->setPluginNamespace(mPluginNamespace.c_str()); return pluginObj; } diff --git a/src/pluginsRT/ReshapeRT.cpp b/src/pluginsRT/ReshapeRT.cpp index 127e244..f364ab0 100644 --- a/src/pluginsRT/ReshapeRT.cpp +++ b/src/pluginsRT/ReshapeRT.cpp @@ -54,10 +54,11 @@ int ReshapeRT::enqueue(int batchSize, const void *const *inputs, void *const *ou } #elif NV_TENSORRT_MAJOR == 7 int32_t ReshapeRT::enqueue(int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) { + std::cout << new_dim.c << ":" << new_dim.h << std::endl; dnnType *srcData = (dnnType*)reinterpret_cast(inputs[0]); dnnType *dstData = reinterpret_cast(outputs[0]); - - checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream)); + std::cout << "C : " << c << "H : " << h << "w :" << w << std::endl; + checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*new_dim.c*new_dim.h*new_dim.w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream)); return 0; } #endif