Adding dilated convolution and support for reorg3d layer

include/tkDNN/DarknetParser.h

@@ -24,6 +24,7 @@ namespace tk { namespace dnn {
         int pad = 0;
         int coords = 4;
         float scale_xy = 1;
+        int dilation = 1;
         std::vector<int> layers;
         std::string activation = "linear";
 

include/tkDNN/Layer.h

@@ -261,6 +261,10 @@ class Conv2d : public LayerWgs {
     Conv2d( Network *net, int out_ch, int kernelH, int kernelW, 
                 int strideH, int strideW, int paddingH, int paddingW,
                 std::string fname_weights, bool batchnorm = false, bool deConv = false, int groups = 1, bool additional_bias=false);
+    Conv2d( Network *net, int out_ch, int kernelH, int kernelW,
+                int strideH, int strideW, int paddingH, int paddingW,
+                int dilationW, int dilationH,
+                std::string fname_weights, bool batchnorm = false, bool deConv = false, int groups = 1, bool additional_bias=false);
     virtual ~Conv2d();
     virtual layerType_t getLayerType() { return LAYER_CONV2D; };
 
@@ -269,6 +273,7 @@ class Conv2d : public LayerWgs {
     int kernelH, kernelW, strideH, strideW, paddingH, paddingW;
     bool deConv, additional_bias;
     int groups;
+    int dilationW, dilationH;
 
 protected:
     cudnnFilterDescriptor_t filterDesc;
@@ -529,13 +534,14 @@ class Route : public Layer {
 class Reorg : public Layer {
 
 public:
-    Reorg(Network *net, int stride);
+    Reorg(Network *net, int stride, bool reorg3d = false);
     virtual ~Reorg();
     virtual layerType_t getLayerType() { return LAYER_REORG; };
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
 
     int stride;
+    bool reorg3d;
 };
 
 /**

include/tkDNN/pluginsRT/ReorgRT.h

@@ -4,8 +4,9 @@
 class ReorgRT : public IPlugin {
 
 public:
-	ReorgRT(int stride) {
+	ReorgRT(int stride, bool reorg3d = false) {
 		this->stride = stride;
+		this->reorg3d = reorg3d;
 	}
 
 	~ReorgRT(){
@@ -21,9 +22,15 @@ class ReorgRT : public IPlugin {
 	}
 
 	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
-		c = inputDims[0].d[0];
-		h = inputDims[0].d[1];
-		w = inputDims[0].d[2];
+		if (reorg3d) {
+			c = outputDims[0].d[0];
+			h = outputDims[0].d[1];
+			w = outputDims[0].d[2];
+		} else {
+			c = inputDims[0].d[0];
+			h = inputDims[0].d[1];
+			w = inputDims[0].d[2];
+		}
 	}
 
 	int initialize() override {
@@ -60,4 +67,5 @@ class ReorgRT : public IPlugin {
 	}
 
 	int c, h, w, stride;
+	bool reorg3d;
 };

src/Conv2d.cpp

@@ -31,9 +31,9 @@ void Conv2d::initCUDNN(bool back) {
                                            kernelH, kernelW) );
 
     checkCUDNN( cudnnSetConvolution2dDescriptor(convDesc,
-                                                paddingH, paddingW, // padding
+                                                paddingH * dilationH, paddingW * dilationW, // padding
                                                 strideH, strideW, // stride
-                                                1,1, // upscale
+                                                dilationH, dilationW, // upscale
                                                 CUDNN_CROSS_CORRELATION, CUDNN_DATA_FLOAT) );
 
     checkCUDNN(  cudnnSetConvolutionGroupCount(convDesc,
@@ -146,6 +146,50 @@ Conv2d::Conv2d( Network *net, int out_ch, int kernelH, int kernelW,
     this->groups = groups;
     this->additional_bias = additional_bias;
 
+    this->dilationW = 1;
+    this->dilationH = 1;
+
+    if(!deConv) {
+        output_dim.n = input_dim.n;
+        output_dim.c = out_ch;
+        output_dim.h = (input_dim.h + 2 * paddingH - kernelH) / strideH + 1;
+        output_dim.w = (input_dim.w + 2 * paddingW - kernelW) / strideW + 1;
+        output_dim.l = 1;
+    } else {
+        output_dim.n = input_dim.n;
+        output_dim.c = out_ch;
+        output_dim.h = ((input_dim.h-1) * strideH) - 2*paddingH + kernelH;
+        output_dim.w = ((input_dim.w-1) * strideW) - 2*paddingW + kernelW;
+        output_dim.l = 1;
+    }
+    initCUDNN(deConv);
+
+    // allocate warkspace
+    if (ws_sizeInBytes!=0) {
+        checkCuda( cudaMalloc(&workSpace, ws_sizeInBytes) );
+    }
+
+    //allocate data for infer result
+    checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) );
+}
+
+Conv2d::Conv2d(Network *net, int out_ch, int kernelH, int kernelW, int strideH, int strideW, int paddingH, int paddingW, int dilationX, int dilationY, std::string fname_weights, bool batchnorm, bool deConv, int groups, bool additional_bias)
+    :LayerWgs(net, net->getOutputDim().c, out_ch, kernelH, kernelW, 1,
+              fname_weights, batchnorm, additional_bias, deConv, groups)
+{
+    this->dilationW = dilationX;
+    this->dilationH = dilationY;
+
+    this->kernelH = kernelH;
+    this->kernelW = kernelW;
+    this->strideH = strideH;
+    this->strideW = strideW;
+    this->paddingH = paddingH;
+    this->paddingW = paddingW;
+    this->deConv = deConv;
+    this->groups = groups;
+    this->additional_bias = additional_bias;
+
     if(!deConv) {
         output_dim.n = input_dim.n;
         output_dim.c = out_ch;

src/DarknetParser.cpp

@@ -83,6 +83,8 @@ namespace tk { namespace dnn {
             auto vec = fromStringToIntVec(value, ',');
             fields.n_mask = vec.size();
         }
+        else if(name.find("dilation") !=  std::string::npos)
+            fields.dilation = std::stoi(value);
         else if(name.find("layers") !=  std::string::npos)
             fields.layers = fromStringToIntVec(value, ',');
 
@@ -111,7 +113,7 @@ namespace tk { namespace dnn {
             std::string wgs = wgs_path + "/c" + std::to_string(netLayers.size()) + ".bin";
             //printf("%d (%d,%d) (%d,%d) (%d,%d) %s %d %d\n", f.filters, f.size_x, f.size_y, f.stride_x, f.stride_y, f.padding_x, f.padding_y, wgs.c_str(), f.batch_normalize, f.groups);
             tk::dnn::Conv2d *l= new tk::dnn::Conv2d(net, f.filters, f.size_x, f.size_y, f.stride_x, 
-                                f.stride_y, f.padding_x, f.padding_y, wgs, f.batch_normalize, false, f.groups);
+                                f.stride_y, f.padding_x, f.padding_y, f.dilation, f.dilation, wgs, f.batch_normalize, false, f.groups);
             netLayers.push_back(l);
         } else if(f.type == "maxpool") {
             if(f.stride_x == 1 && f.stride_y == 1)
@@ -150,8 +152,8 @@ namespace tk { namespace dnn {
             }
             netLayers.push_back(new tk::dnn::Route(net, layers.data(), layers.size()));
 
-        } else if(f.type == "reorg") {
-            netLayers.push_back(new tk::dnn::Reorg(net, f.stride_x));
+        } else if(f.type == "reorg" || f.type == "reorg3d") {
+            netLayers.push_back(new tk::dnn::Reorg(net, f.stride_x, f.type == "reorg3d"));
 
         } else if(f.type == "region") {
             netLayers.push_back(new tk::dnn::Region(net, f.classes, f.coords, f.num));

src/NetworkRT.cpp

@@ -316,7 +316,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
             l->outputs, DimsHW{l->kernelH, l->kernelW}, w, b);
         checkNULL(lRTconv);
         lRTconv->setStride(DimsHW{l->strideH, l->strideW});
-        lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
+        lRTconv->setPadding(DimsHW{l->paddingH * l->dilationH, l->paddingW * l->dilationW});
+        lRTconv->setDilation(DimsHW{l->dilationH, l->dilationW});
         lRTconv->setNbGroups(l->groups);
         lRT = (ILayer*) lRTconv;
     } else {
@@ -480,7 +481,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Reorg *l) {
     //std::cout<<"convert Reorg\n";
 
     //std::cout<<"New plugin REORG\n";
-    IPlugin *plugin = new ReorgRT(l->stride);
+    IPlugin *plugin = new ReorgRT(l->stride, l->reorg3d);
     IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
     checkNULL(lRT);
     return lRT;

src/Reorg.cpp

@@ -5,9 +5,10 @@
 
 namespace tk { namespace dnn {
 
-Reorg::Reorg(Network *net, int stride) : Layer(net) {
+Reorg::Reorg(Network *net, int stride, bool reorg3d) : Layer(net) {
 
     this->stride = stride;
+    this->reorg3d = reorg3d;
 
     output_dim.n = input_dim.n;
     output_dim.c = input_dim.c*stride*stride;
@@ -24,8 +25,10 @@ Reorg::~Reorg() {
 }
 
 dnnType* Reorg::infer(dataDim_t &dim, dnnType* srcData) {
-
-    reorgForward(srcData, dstData, dim.n, dim.c, dim.h, dim.w, stride);
+    if (reorg3d)
+        reorgForward(srcData, dstData, output_dim.n, output_dim.c, output_dim.h, output_dim.w, stride);
+    else
+        reorgForward(srcData, dstData, dim.n, dim.c, dim.h, dim.w, stride);
 
     dim = output_dim;
     return dstData;

tests/darknet/yolo3_tiny_dilation.cpp

@@ -0,0 +1,33 @@
+#include<iostream>
+#include<vector>
+#include "tkdnn.h"
+#include "test.h"
+#include "DarknetParser.h"
+
+int main() {
+    std::string bin_path  = "dilation";
+    std::vector<std::string> input_bins = {
+        bin_path + "/layers/input.bin"
+    };
+    std::vector<std::string> output_bins = {
+        bin_path + "/debug/layer33_out.bin",
+        bin_path + "/debug/layer45_out.bin",
+    };
+    std::string wgs_path  = bin_path + "/layers";
+    std::string cfg_path  = "config.cfg";
+    std::string name_path = "config.names";
+    //downloadWeightsifDoNotExist(input_bins[0], bin_path, "https://cloud.hipert.unimore.it/s/LMcSHtWaLeps8yN/download");
+
+    // parse darknet network
+    tk::dnn::Network *net = tk::dnn::darknetParser(cfg_path, wgs_path, name_path);
+    net->print();
+
+    //convert network to tensorRT
+    tk::dnn::NetworkRT *netRT = new tk::dnn::NetworkRT(net, net->getNetworkRTName(bin_path.c_str()));
+
+    int ret = testInference(input_bins, output_bins, net, netRT);
+    net->releaseLayers();
+    delete net;
+    delete netRT;
+    return ret;
+}