From 7c81c5a43cdc04550d1c65260c1394551307d21f Mon Sep 17 00:00:00 2001 From: Francesco Gatti Date: Tue, 21 Apr 2020 19:20:44 +0200 Subject: [PATCH] batch size > 1 --- include/tkDNN/Network.h | 1 + include/tkDNN/NetworkRT.h | 9 ++++- .../tkDNN/pluginsRT/MaxPoolingFixedSizeRT.h | 2 +- src/Network.cpp | 4 +++ src/NetworkRT.cpp | 24 +++++++++----- tests/test_rtinference/rtinference.cpp | 33 ++++++++++++------- 6 files changed, 51 insertions(+), 22 deletions(-) diff --git a/include/tkDNN/Network.h b/include/tkDNN/Network.h index eb17ea8..0e8fef8 100644 --- a/include/tkDNN/Network.h +++ b/include/tkDNN/Network.h @@ -62,6 +62,7 @@ public: dataDim_t getOutputDim(); bool fp16, dla, int8; + int maxBatchSize; bool dontLoadWeights; std::string fileImgList; std::string fileLabelList; diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h index 64deac9..1152234 100644 --- a/include/tkDNN/NetworkRT.h +++ b/include/tkDNN/NetworkRT.h @@ -74,11 +74,18 @@ public: NetworkRT(Network *net, const char *name); virtual ~NetworkRT(); + int getMaxBatchSize() { + if(engineRT != nullptr) + return engineRT->getMaxBatchSize(); + else + return 0; + } + /** Do inferece */ dnnType* infer(dataDim_t &dim, dnnType* data); - void enqueue(); + void enqueue(int batchSize = 1); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Layer *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Conv2d *l); diff --git a/include/tkDNN/pluginsRT/MaxPoolingFixedSizeRT.h b/include/tkDNN/pluginsRT/MaxPoolingFixedSizeRT.h index 992042b..2efa6e2 100644 --- a/include/tkDNN/pluginsRT/MaxPoolingFixedSizeRT.h +++ b/include/tkDNN/pluginsRT/MaxPoolingFixedSizeRT.h @@ -42,7 +42,7 @@ public: virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override { - std::cout<n<<" "<c<<" "<h<<" "<w<<" "<stride_H<<" "<stride_W<<" "<winSize<<" "<padding<n<<" "<c<<" "<h<<" "<w<<" "<stride_H<<" "<stride_W<<" "<winSize<<" "<padding<(inputs[0]); dnnType *dstData = reinterpret_cast(outputs[0]); MaxPoolingForward(srcData, dstData, this->n, this->c, this->h, this->w, this->stride_H, this->stride_W, this->winSize, this->padding); diff --git a/src/Network.cpp b/src/Network.cpp index 885ec1a..f199fc2 100644 --- a/src/Network.cpp +++ b/src/Network.cpp @@ -34,6 +34,10 @@ Network::Network(dataDim_t input_dim) { int8 = true; } } + maxBatchSize = 1; + if(const char* env_p = std::getenv("TKDNN_BATCHSIZE")) { + maxBatchSize = atoi(env_p); + } if(const char* env_p = std::getenv("TKDNN_CALIB_IMG_PATH")) fileImgList = env_p; diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 6a8dabe..6ed6efe 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -58,7 +58,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) { dataDim_t dim = net->layers[0]->input_dim; dtRT = DataType::kFLOAT; - builderRT->setMaxBatchSize(1); + builderRT->setMaxBatchSize(net->maxBatchSize); builderRT->setMaxWorkspaceSize(1 << 30); if(net->fp16 && builderRT->platformHasFastFp16()) { @@ -133,6 +133,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) { input->setName("out"); networkRT->markOutput(*input); + std::cout<<"Selected maxBatchSize: "<getMaxBatchSize()<<"\n"; std::cout<<"Building tensorRT cuda engine...\n"; #if NV_TENSORRT_MAJOR >= 6 engineRT = builderRT->buildEngineWithConfig(*networkRT, *configRT); @@ -181,9 +182,9 @@ NetworkRT::NetworkRT(Network *net, const char *name) { // create GPU buffers and a stream for(int i=0; igetNbBindings(); i++) { Dims dim = engineRT->getBindingDimensions(i); - checkCuda(cudaMalloc(&buffersRT[i], dim.d[0]*dim.d[1]*dim.d[2]*sizeof(dnnType))); + checkCuda(cudaMalloc(&buffersRT[i], engineRT->getMaxBatchSize()*dim.d[0]*dim.d[1]*dim.d[2]*sizeof(dnnType))); } - checkCuda(cudaMalloc(&output, output_dim.tot()*sizeof(dnnType))); + checkCuda(cudaMalloc(&output, engineRT->getMaxBatchSize()*output_dim.tot()*sizeof(dnnType))); checkCuda(cudaStreamCreate(&stream)); } @@ -192,19 +193,24 @@ NetworkRT::~NetworkRT() { } dnnType* NetworkRT::infer(dataDim_t &dim, dnnType* data) { + int batches = dim.n; + if(batches > getMaxBatchSize()) { + FatalError("input batch size too large"); + } - checkCuda(cudaMemcpyAsync(buffersRT[buf_input_idx], data, input_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream)); - contextRT->enqueue(1, buffersRT, stream, nullptr); - checkCuda(cudaMemcpyAsync(output, buffersRT[buf_output_idx], output_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream)); - cudaStreamSynchronize(stream); + checkCuda(cudaMemcpyAsync(buffersRT[buf_input_idx], data, batches*input_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream)); + contextRT->enqueue(batches, buffersRT, stream, nullptr); + checkCuda(cudaMemcpyAsync(output, buffersRT[buf_output_idx], batches*output_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream)); + checkCuda(cudaStreamSynchronize(stream)); dim = output_dim; + dim.n = batches; return output; } -void NetworkRT::enqueue() { - contextRT->enqueue(1, buffersRT, stream, nullptr); +void NetworkRT::enqueue(int batchSize) { + contextRT->enqueue(batchSize, buffersRT, stream, nullptr); } ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) { diff --git a/tests/test_rtinference/rtinference.cpp b/tests/test_rtinference/rtinference.cpp index f99d003..88aedf0 100644 --- a/tests/test_rtinference/rtinference.cpp +++ b/tests/test_rtinference/rtinference.cpp @@ -2,31 +2,42 @@ #include "tkdnn.h" #include /* srand, rand */ + int main(int argc, char *argv[]) { if(argc < 2 || !fileExist(argv[1])) FatalError("unable to read serialRT file"); + int BATCH_SIZE = 1; + if(argc >2) + BATCH_SIZE = atoi(argv[2]); + //always same test srand (0); //convert network to tensorRT tk::dnn::NetworkRT netRT(NULL, argv[1]); + + tk::dnn::dataDim_t idim = netRT.input_dim; + tk::dnn::dataDim_t odim = netRT.output_dim; + idim.n = BATCH_SIZE; + odim.n = BATCH_SIZE; + dnnType *input = new float[idim.tot()]; + dnnType *output = new float[odim.tot()]; + dnnType *input_d; + checkCuda( cudaMalloc(&input_d, idim.tot()*sizeof(dnnType))); - dnnType *input = new float[netRT.input_dim.tot()]; - dnnType *output = new float[netRT.input_dim.tot()]; - + std::cout<<"Testing with batchsize: "<