Yolo3_tiny CUDNN works, TensorRT doesn't. Add n_masks to Yolo layer.

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

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;
 
@@ -253,14 +253,18 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
         lRTconv->setStride(DimsHW{l->strideH, l->strideW});
         lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
         lRT = (ILayer*) lRTconv;
+        
     }
 
+    checkNULL(lRT);
     if(l->batchnorm) {
         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;
         IScaleLayer *lRT2 = networkRT->addScale(*lRT->getOutput(0), ScaleMode::kCHANNEL, 
                     shift, scale, power);
+        
         checkNULL(lRT2);
 
         Weights shift2{dtRT, bias_b, l->outputs};
@@ -277,7 +281,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
     std::cout<<"convert Pooling\n";
-    // printf("%d %d\n", l->winW, l->winH);
+    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;
@@ -287,8 +291,27 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
     IPoolingLayer *lRT = networkRT->addPooling(*input, 
         ptype, DimsHW{l->winH, l->winW});
     checkNULL(lRT);
-    lRT->setStride(DimsHW{l->strideH, l->strideW});
-    lRT->setPadding(DimsHW{l->paddingH, l->paddingW});
+    
+    // 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++) {
+            std::cout<<t->getDimensions().d[j]<<" ";
+        }
+        std::cout<<" (TensorRT)\n";
 
     return lRT;
 }
@@ -324,12 +347,19 @@ 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++) {
+            std::cout<<tens[i]->getDimensions().d[j]<<" ";
+        }
+        std::cout<<"\n";
     }
+    
     IConcatenationLayer *lRT = networkRT->addConcatenation(tens, l->layers_n);
     //IPlugin *plugin = new RouteRT();
     //IPluginLayer *lRT = networkRT->addPlugin(tens, l->layers_n, *plugin);
@@ -474,13 +504,15 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa
 
     if(name.find("Yolo") == 0) {
         YoloRT *r = new YoloRT(readBUF<int>(buf),    //classes
-                                readBUF<int>(buf));   //num
+                                readBUF<int>(buf),   //num
+                                nullptr,
+                                readBUF<int>(buf));   //n_masks
         r->c = readBUF<int>(buf);
         r->h = readBUF<int>(buf);
         r->w = readBUF<int>(buf);
-        for(int i=0; i<r->num; i++)
+        for(int i=0; i<r->n_masks; i++)
             r->mask[i] = readBUF<dnnType>(buf);
-        for(int i=0; i<3*2*r->num; i++)
+        for(int i=0; i<r->n_masks*2*r->num; i++)
             r->bias[i] = readBUF<dnnType>(buf);
 
 		// save classes names

src/Yolo.cpp

@@ -11,19 +11,20 @@
 
 namespace tk { namespace dnn {
 
-Yolo::Yolo(Network *net, int classes, int num, std::string fname_weights) : 
+Yolo::Yolo(Network *net, int classes, int num, std::string fname_weights, int n_masks) : 
     Layer(net) {
     
     this->classes = classes;
-    this->num = 3;
+    this->num = num;
+    this->n_masks = n_masks;
 
     // load anchors
     if(fname_weights != "") {
         int seek = 0;
-        readBinaryFile(fname_weights, 3, &mask_h, &mask_d, seek);
-        seek += 3;
-        readBinaryFile(fname_weights, 3*num*2, &bias_h, &bias_d, seek);
-        for(int i=0; i<3*num*2; i++)
+        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]);
     }
 
@@ -70,7 +71,7 @@ dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) {
     checkCuda( cudaMemcpy(dstData, srcData, dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice));
 
     for (int b = 0; b < dim.n; ++b){
-        for(int n = 0; n < num; ++n){
+        for(int n = 0; n < n_masks; ++n){
             int index = entry_index(b, n*dim.w*dim.h, 0, classes, input_dim, output_dim);
             activationLOGISTICForward(srcData + index, dstData + index, 2*dim.w*dim.h);