From 7e7f480e2ba1856a9004477a203bdf6ee8274532 Mon Sep 17 00:00:00 2001
From: fbagni <gattinomicino>
Date: Mon, 23 Dec 2019 16:11:41 +0100
Subject: [PATCH] Yolov3_tiny works on tensorRT

Signed-off-by: fbagni <gattinomicino>
---
 src/Conv2d.cpp        |  4 ++--
 src/NetworkRT.cpp     | 36 ++++++++++++++----------------------
 src/Pooling.cpp       |  4 ----
 src/Yolo.cpp          |  4 ++--
 src/kernels/resize.cu |  4 ++--
 5 files changed, 20 insertions(+), 32 deletions(-)

diff --git a/src/Conv2d.cpp b/src/Conv2d.cpp
index b8296ac..ef73e40 100644
--- a/src/Conv2d.cpp
+++ b/src/Conv2d.cpp
@@ -17,8 +17,8 @@ void Conv2d::initCUDNN(bool back) {
         idim = output_dim;
         odim = input_dim;
     }
-    idim.print();
-    odim.print();
+    //idim.print();
+    //odim.print();
 
     checkCUDNN( cudnnCreateFilterDescriptor(&filterDesc) );
     checkCUDNN( cudnnCreateConvolutionDescriptor(&convDesc) );
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index ed72c89..c744b86 100644
--- a/src/NetworkRT.cpp
+++ b/src/NetworkRT.cpp
@@ -291,31 +291,23 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
     if(l->pool_mode == tkdnnPoolingMode_t::POOLING_AVERAGE) ptype = PoolingType::kAVERAGE;
     if(l->pool_mode == tkdnnPoolingMode_t::POOLING_AVERAGE_EXCLUDE_PADDING) ptype = PoolingType::kMAX_AVERAGE_BLEND;
 
-    IPoolingLayer *lRT = networkRT->addPooling(*input, 
-        ptype, DimsHW{l->winH, l->winW});
+
+    if(l->input_dim.h % 2 == 1 && l->input_dim.w % 2 == 1)
+    {
+      IPlugin *plugin = new ResizeLayerRT( l->output_dim.c,l->output_dim.h+1,l->output_dim.w+1 );
+      IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
+      checkNULL(lRT);
+      lRT->setName( "Resize" );
+
+      input = lRT->getOutput(0);
+    }
+
+    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});
-
-    
-    ITensor *t = lRT->getOutput(0);
-    // for(int j=0; j<t->getDimensions().nbDims; j++) {
-    //         std::cout<<t->getDimensions().d[j]<<" ";
-    //     }
-    //     std::cout<<" (TensorRT)\n";
-
-    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;
+    return lRT;
 }
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {
@@ -551,7 +543,7 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa
         return r;
     } 
 
-    if(name.find("Pooling") == 0) {
+    if(name.find("Resize") == 0) {
         ResizeLayerRT *r = new ResizeLayerRT(readBUF<int>(buf), //o_c
                                             readBUF<int>(buf), //o_h
                                             readBUF<int>(buf)); //o_w
diff --git a/src/Pooling.cpp b/src/Pooling.cpp
index 0f3573d..236adbf 100644
--- a/src/Pooling.cpp
+++ b/src/Pooling.cpp
@@ -25,8 +25,6 @@ Pooling::Pooling( Network *net, int winH, int winW, int strideH, int strideW,
     int h = input_dim.h;
     int w = input_dim.w;
     int l = input_dim.l;
-
-    printf("before: %d %d\n", h, w);
     
 
     poolOn3d = false;
@@ -64,8 +62,6 @@ Pooling::Pooling( Network *net, int winH, int winW, int strideH, int strideW,
 
     checkCUDNN( cudnnSetTensor4dDescriptor(dstTensorDesc,
                 net->tensorFormat, net->dataType, n, c, h, w) );
-
-    printf("after: %d %d\n", h, w);
        
     output_dim.n = n;
     output_dim.c = c;
diff --git a/src/Yolo.cpp b/src/Yolo.cpp
index ca6ebc5..634f5ff 100644
--- a/src/Yolo.cpp
+++ b/src/Yolo.cpp
@@ -24,8 +24,8 @@ Yolo::Yolo(Network *net, int classes, int num, std::string fname_weights, int n_
         readBinaryFile(fname_weights, n_masks, &mask_h, &mask_d, seek);
         seek += n_masks;
         readBinaryFile(fname_weights, n_masks*num*2, &bias_h, &bias_d, seek);
-        for(int i=0; i<n_masks*num*2; i++)
-            printf("%f\n", bias_h[i]);
+        //for(int i=0; i<n_masks*num*2; i++)
+            //printf("%f\n", bias_h[i]);
     }
 
     // init default classes name
diff --git a/src/kernels/resize.cu b/src/kernels/resize.cu
index f637152..169b853 100644
--- a/src/kernels/resize.cu
+++ b/src/kernels/resize.cu
@@ -17,7 +17,7 @@ __global__ void resize_kernel(  int i_N,float *x, int i_w, int i_h, int i_c,
     int out_c = i%o_c;
     i = i/o_c;
 
-    //copying last column/last row
+    //copying last column/last row as padding
     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];
 }
@@ -43,4 +43,4 @@ void resizeForward(  dnnType* srcData, dnnType* dstData, int n, int i_c, int i_h
         // printDeviceVector(i_size, srcData);
         // printDeviceVector(o_size, dstData);
     }
-}
\ No newline at end of file
+}