Add ResizeLayerRT plugin

Signed-off-by: nvidia <micaelaverucchi@gmail.com>

demo/demo/demo.cpp

@@ -25,7 +25,7 @@ int main(int argc, char *argv[]) {
     signal(SIGINT, sig_handler);
 
 
-    char *net = "yolo3_berkeley.rt";
+    char *net = "yolo3.rt";
     if(argc > 1)
         net = argv[1]; 
     char *input = "../demo/yolo_test.mp4";

include/tkDNN/NetworkRT.h

@@ -30,6 +30,7 @@ using namespace nvinfer1;
 #include "pluginsRT/ShortcutRT.h"
 #include "pluginsRT/YoloRT.h"
 #include "pluginsRT/UpsampleRT.h"
+#include "pluginsRT/ResizeLayerRT.h"
 //#include "pluginsRT/Int8Calibrator.h"
 
 class PluginFactory : IPluginFactory

include/tkDNN/Yolo3Detection.h

@@ -33,6 +33,7 @@ class Yolo3Detection {
     public:
         int classes = 0;
         int num = 0;
+        int n_masks = 0;
         float thresh = 0.3;
         cv::Scalar colors[256];
 

include/tkDNN/kernels.h

@@ -9,6 +9,9 @@ void activationLOGISTICForward(dnnType* srcData, dnnType* dstData, int size, cud
 
 void fill(dnnType* data, int size, dnnType val, cudaStream_t stream = cudaStream_t(0));
 
+void resizeForward(  dnnType* srcData, dnnType* dstData, int n, int i_c, int i_h, int i_w,
+                    int o_c, int o_h, int o_w, cudaStream_t stream = cudaStream_t(0));
+
 void reorgForward(  dnnType* srcData, dnnType* dstData, 
                     int n, int c, int h, int w, int stride, cudaStream_t stream = cudaStream_t(0));
 void softmaxForward(float *input, int n, int batch, int batch_offset, 

include/tkDNN/pluginsRT/ResizeLayerRT.h

@@ -0,0 +1,67 @@
+#include<cassert>
+#include "../kernels.h"
+
+class ResizeLayerRT : public IPlugin {
+
+public:
+	ResizeLayerRT(int c, int h, int w) {
+		o_c = c;
+		o_h = h;
+		o_w = w;	
+	}
+
+	~ResizeLayerRT(){
+	}
+
+	int getNbOutputs() const override {
+		return 1;
+	}
+
+	Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
+		return DimsCHW{o_c, o_h, o_w};
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
+		i_c = inputDims[0].d[0];
+		i_h = inputDims[0].d[1];
+		i_w = inputDims[0].d[2];        
+	}
+
+	int initialize() override {
+		return 0;
+	}
+
+	virtual void terminate() override {
+	}
+
+	virtual size_t getWorkspaceSize(int maxBatchSize) const override {
+		return 0;
+	}
+
+	virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
+    	// printf("%d %d %d %d %d %d\n", i_c, i_w, i_h, o_c, o_w, o_h);
+        resizeForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]), 
+					  reinterpret_cast<dnnType*>(outputs[0]), 
+					  batchSize, i_c, i_h, i_w, o_c, o_h, o_w, stream);
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 6*sizeof(int);
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer);
+
+		tk::dnn::writeBUF(buf, o_c);
+		tk::dnn::writeBUF(buf, o_h);
+		tk::dnn::writeBUF(buf, o_w);
+
+		tk::dnn::writeBUF(buf, i_c);
+		tk::dnn::writeBUF(buf, i_h);
+		tk::dnn::writeBUF(buf, i_w);
+	}
+
+	int i_c, i_h, i_w, o_c, o_h, o_w;
+};

src/NetworkRT.cpp

@@ -15,9 +15,9 @@ using namespace nvinfer1;
 // Logger for info/warning/errors
 class Logger : public ILogger {
     void log(Severity severity, const char* msg) override {
-// #ifdef DEBUG
+#ifdef DEBUG
         std::cout <<"TENSORRT LOG: "<< msg << std::endl;
-// #endif
+#endif
     }
 } loggerRT;
 
@@ -209,8 +209,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Dense *l) {
 
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
-    std::cout<<"convert conv2D\n";
+    // std::cout<<"convert conv2D\n";
-    printf("%d %d %d %d %d\n", l->kernelH, l->kernelW, l->inputs, l->outputs, l->batchnorm);
+    // printf("%d %d %d %d %d\n", l->kernelH, l->kernelW, l->inputs, l->outputs, l->batchnorm);
 
 
     void *data_b, *bias_b, *power_b, *mean_b, *variance_b, *scales_b;
@@ -261,7 +261,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
         Weights power{dtRT, power_b, l->outputs};
         Weights shift{dtRT, mean_b, l->outputs};
         Weights scale{dtRT, variance_b, l->outputs};
-        std::cout<<lRT->getNbOutputs()<<std::endl;
+        // std::cout<<lRT->getNbOutputs()<<std::endl;
         IScaleLayer *lRT2 = networkRT->addScale(*lRT->getOutput(0), ScaleMode::kCHANNEL, 
                     shift, scale, power);
         
@@ -280,8 +280,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
 }
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
-    std::cout<<"convert Pooling\n";
+    // std::cout<<"convert Pooling\n";
-    printf("%d %d %d %d %d %d %d %d %d %d %d %d (layer)\n", l->input_dim.h, l->input_dim.w, l->output_dim.h, l->output_dim.w,  l->winW, l->winH,  l->strideH, l->strideW, l->paddingH, l->paddingW, l->pool_mode, tkdnnPoolingMode_t::POOLING_MAX) ;
 
     PoolingType ptype;
     if(l->pool_mode == tkdnnPoolingMode_t::POOLING_MAX) ptype = PoolingType::kMAX;
@@ -291,29 +290,28 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
     IPoolingLayer *lRT = networkRT->addPooling(*input, 
         ptype, DimsHW{l->winH, l->winW});
     checkNULL(lRT);
+
+    lRT->setPadding(DimsHW{l->paddingH, l->paddingW});
+    lRT->setStride(DimsHW{l->strideH, l->strideW});
+
     
-    // if (l->input_dim.h == 13 && l->output_dim.h == 13)
-    // {
-    //     lRT->setPadding(DimsHW{7, 7});    
-    //     lRT->setStride(DimsHW{2, 2});
-    // }
-    // else
-    // {
-        lRT->setPadding(DimsHW{l->paddingH, l->paddingW});
-        lRT->setStride(DimsHW{l->strideH, l->strideW});
-    // }
-
-    // IResizeLayer *lRT = networkRT->addResize(*lRT->getOutput(0));
-    // checkNULL(lRT);
-    // lRT->setOutputDimensions(l->output_dim);
-
     ITensor *t = lRT->getOutput(0);
-    for(int j=0; j<t->getDimensions().nbDims; j++) {
+    // for(int j=0; j<t->getDimensions().nbDims; j++) {
-            std::cout<<t->getDimensions().d[j]<<" ";
+    //         std::cout<<t->getDimensions().d[j]<<" ";
-        }
+    //     }
-        std::cout<<" (TensorRT)\n";
+    //     std::cout<<" (TensorRT)\n";
 
-    return lRT;
+    IPlugin *plugin = new ResizeLayerRT( l->output_dim.c,l->output_dim.h,l->output_dim.w );
+    IPluginLayer *lRT1 = networkRT->addPlugin(&t, 1, *plugin);
+    checkNULL(lRT1);
+
+    // ITensor *t1 = lRT1->getOutput(0);
+    // for(int j=0; j<t1->getDimensions().nbDims; j++) {
+    //         std::cout<<t1->getDimensions().d[j]<<" ";
+    //     }
+    //     std::cout<<" (TensorRT after resize )\n";
+
+    return lRT1;
 }
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {
@@ -347,17 +345,17 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Softmax *l) {
 }
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Route *l) {
-    std::cout<<"convert route\n";
+    // std::cout<<"convert route\n";
     
 
 
     ITensor **tens = new ITensor*[l->layers_n];
     for(int i=0; i<l->layers_n; i++) {
         tens[i] = tensors[l->layers[i]];
-        for(int j=0; j<tens[i]->getDimensions().nbDims; j++) {
+        // for(int j=0; j<tens[i]->getDimensions().nbDims; j++) {
-            std::cout<<tens[i]->getDimensions().d[j]<<" ";
+        //     std::cout<<tens[i]->getDimensions().d[j]<<" ";
-        }
+        // }
-        std::cout<<"\n";
+        // std::cout<<"\n";
     }
     
     IConcatenationLayer *lRT = networkRT->addConcatenation(tens, l->layers_n);
@@ -502,6 +500,16 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa
         return r;
     } 
 
+    if(name.find("Pooling") == 0) {
+        ResizeLayerRT *r = new ResizeLayerRT(readBUF<int>(buf), //o_c
+                                            readBUF<int>(buf), //o_h
+                                            readBUF<int>(buf)); //o_w
+        r->i_c = readBUF<int>(buf);
+        r->i_h = readBUF<int>(buf);
+        r->i_w = readBUF<int>(buf);
+        return r;
+    } 
+
     if(name.find("Yolo") == 0) {
         YoloRT *r = new YoloRT(readBUF<int>(buf),    //classes
                                 readBUF<int>(buf),   //num

src/Pooling.cpp

@@ -53,7 +53,6 @@ Pooling::Pooling( Network *net, int winH, int winW, int strideH, int strideW,
     // checkCUDNN( cudnnGetPooling2dForwardOutputDim(poolingDesc, srcTensorDesc, &n, &c, &h, &w)); 
 
     //compute w and h as in darknet
-
     int padH = paddingH == 0? winH -1 : paddingH;
     int padW = paddingW == 0? winW -1 : paddingW;
 

src/Yolo.cpp

@@ -130,7 +130,7 @@ int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int netw, int net
     for (i = 0; i < lw*lh; ++i){
         int row = i / lw;
         int col = i % lw;
-        for(n = 0; n < num; ++n){
+        for(n = 0; n < n_masks; ++n){
             int obj_index  = entry_index(0, n*lw*lh + i, 4, classes, input_dim, output_dim);
             float objectness = predictions[obj_index];
             if(objectness <= thresh) continue;

src/Yolo3Detection.cpp

@@ -22,7 +22,8 @@ bool Yolo3Detection::init(std::string tensor_path) {
     std::cout<<(tensor_path).c_str()<<"\n";
     netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
 
-    if(netRT->pluginFactory->n_yolos != 3) {
+
+    if(netRT->pluginFactory->n_yolos < 2 ) {
         FatalError("this is not yolo3");
     }
 
@@ -30,13 +31,14 @@ bool Yolo3Detection::init(std::string tensor_path) {
         YoloRT *yRT = netRT->pluginFactory->yolos[i];
         classes = yRT->classes;
         num = yRT->num;
+        n_masks = yRT->n_masks;
 
         // make a yolo layer for interpret predictions
-        yolo[i] = new tk::dnn::Yolo(nullptr, classes, num, ""); // yolo without input and bias
+        yolo[i] = new tk::dnn::Yolo(nullptr, classes, n_masks, ""); // yolo without input and bias
-        yolo[i]->mask_h = new dnnType[num];
+        yolo[i]->mask_h = new dnnType[n_masks];
-        yolo[i]->bias_h = new dnnType[num*3*2];
+        yolo[i]->bias_h = new dnnType[num*n_masks*2];
-        memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*num);
+        memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*n_masks);
-        memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*3*2);
+        memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*n_masks*2);
         yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, yRT->c, yRT->h, yRT->w);
         yolo[i]->classesNames = yRT->classesNames;
     }
@@ -84,7 +86,7 @@ void Yolo3Detection::update(cv::Mat &imageORIG) {
 
 
     //DO INFERENCE
-    dnnType *rt_out[3]; 
+    dnnType *rt_out[netRT->pluginFactory->n_yolos]; 
     tk::dnn::dataDim_t dim = netRT->input_dim;
     checkCuda(cudaMemcpyAsync(input_d, input, dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
 
@@ -99,7 +101,7 @@ void Yolo3Detection::update(cv::Mat &imageORIG) {
     TIMER_START
     // compute dets
     ndets = 0;
-    for(int i=0; i<3; i++) {
+    for(int i=0; i<netRT->pluginFactory->n_yolos; i++) {
         rt_out[i] = (dnnType*)netRT->buffersRT[i+1];
         yolo[i]->dstData = rt_out[i];
         yolo[i]->computeDetections(dets, ndets, netRT->input_dim.w, netRT->input_dim.h, thresh);

src/kernels/resize.cu

@@ -0,0 +1,46 @@
+#include "kernels.h"
+#include <stdio.h>
+#define MIN(a,b) (((a)<(b))?(a):(b))
+#define MAX(a,b) (((a)>(b))?(a):(b))
+
+__global__ void resize_kernel(  int i_N,float *x, int i_w, int i_h, int i_c,  
+                                int o_w, int o_h, int o_c, int batch, float *out)
+{
+    int i = (blockIdx.x + blockIdx.y*gridDim.x) * blockDim.x + threadIdx.x;
+    if(i >= i_N) return;
+
+    int out_index = i;
+    int out_w = i%o_w;
+    i = i/o_w;
+    int out_h = i%o_h;
+    i = i/o_h;
+    int out_c = i%o_c;
+    i = i/o_c;
+
+    //copying last column/last row
+    int in_index = ((i*i_c + MIN(out_c,i_c-1))*i_h +  MIN(out_h,i_h-1))*i_w + MIN(out_w, i_w-1);
+    out[out_index] = x[in_index];
+}
+
+
+void resizeForward(  dnnType* srcData, dnnType* dstData, int n, int i_c, int i_h, int i_w,
+     int o_c, int o_h, int o_w, cudaStream_t stream )
+{
+    int i_size = n*i_c*i_h*i_w;
+    int o_size = n*o_c*o_h*o_w;
+
+    int blocks = (o_size+255)/256;
+    int threads = 256;
+
+    if(i_c == o_c && i_h == o_h && i_w == o_w )
+    {
+        checkCuda(cudaMemcpy(dstData, srcData, i_size*sizeof(dnnType), cudaMemcpyDeviceToDevice));
+    }
+    else
+    {
+        checkCuda(cudaMemset(dstData, 0, o_size*sizeof(dnnType)));
+        resize_kernel<<<blocks, threads, 0, stream>>>(o_size, srcData, i_w, i_h, i_c, o_w, o_h, o_c, n, dstData);
+        // printDeviceVector(i_size, srcData);
+        // printDeviceVector(o_size, dstData);
+    }
+}

tests/yolo3_tiny/yolo3_tiny.cpp

@@ -118,5 +118,6 @@ int main() {
     std::cout<<"CUDNN vs correct"; checkResult(out_dim, out_data, out);
     std::cout<<"TRT   vs correct"; checkResult(out_dim, out_data2, out);
     std::cout<<"CUDNN vs TRT    "; checkResult(out_dim, out_data, out_data2);
+    
     return 0;
 }