batch size > 1

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;

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);

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<<this->n<<"  "<<this->c<<"  "<<this->h<<"  "<<this->w<<"  "<<this->stride_H<<"  "<<this->stride_W<<"  "<<this->winSize<<"  "<<this->padding<<std::endl;
+		//std::cout<<this->n<<"  "<<this->c<<"  "<<this->h<<"  "<<this->w<<"  "<<this->stride_H<<"  "<<this->stride_W<<"  "<<this->winSize<<"  "<<this->padding<<std::endl;
 		dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
 		dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
 		MaxPoolingForward(srcData, dstData, this->n, this->c, this->h, this->w, this->stride_H, this->stride_W, this->winSize, this->padding);

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;
     

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: "<<builderRT->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; i<engineRT->getNbBindings(); 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));
+    checkCuda(cudaMemcpyAsync(buffersRT[buf_input_idx], data, batches*input_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
-    contextRT->enqueue(1, buffersRT, stream, nullptr);
+    contextRT->enqueue(batches, buffersRT, stream, nullptr);
-    checkCuda(cudaMemcpyAsync(output, buffersRT[buf_output_idx], output_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
+    checkCuda(cudaMemcpyAsync(output, buffersRT[buf_output_idx], batches*output_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
-    cudaStreamSynchronize(stream);
+    checkCuda(cudaStreamSynchronize(stream));
 
     dim = output_dim;
+    dim.n = batches;
 
     return output;
 }
 
-void NetworkRT::enqueue() {
+void NetworkRT::enqueue(int batchSize) {
-    contextRT->enqueue(1, buffersRT, stream, nullptr);
+    contextRT->enqueue(batchSize, buffersRT, stream, nullptr);
 }
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {

tests/test_rtinference/rtinference.cpp

@@ -2,31 +2,42 @@
 #include "tkdnn.h"
 #include <stdlib.h>     /* 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()];
+    std::cout<<"Testing with batchsize: "<<BATCH_SIZE<<"\n";
-    dnnType *output = new float[netRT.input_dim.tot()];
-
     printCenteredTitle(" TENSORRT inference ", '=', 30); 
-    for(int i=0; i<100; i++) {
+    for(int i=0; i<10; i++) {
-        for(int j=0; j<netRT.input_dim.tot(); j++)
+        for(int j=0; j<idim.tot(); j++) {
             input[j] = ((float) rand() / (RAND_MAX));
+        }
+        checkCuda(cudaMemcpy(input_d, input, idim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice));
+
+        tk::dnn::dataDim_t dim = idim;
         TIMER_START
-        checkCuda(  cudaMemcpyAsync(netRT.buffersRT[netRT.buf_input_idx], input, 
+        netRT.infer(dim, input_d);
-                    netRT.input_dim.tot()*sizeof(float), cudaMemcpyHostToDevice, netRT.stream));
-        netRT.enqueue();
-        checkCuda(  cudaMemcpyAsync(output, netRT.buffersRT[netRT.buf_output_idx], 
-                                    netRT.output_dim.tot()*sizeof(float), cudaMemcpyDeviceToHost, netRT.stream));
-        cudaStreamSynchronize(netRT.stream);
         TIMER_STOP
     }