Add grouped convolutions in CUDNN and tensorRT.

Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>

include/tkDNN/Layer.h

@@ -85,7 +85,7 @@ class LayerWgs : public Layer {
 
 public:
     LayerWgs(Network *net, int inputs, int outputs, int kh, int kw, int kt,
-             std::string fname_weights, bool batchnorm = false, bool additional_bias = false, bool final = false); 
+             std::string fname_weights, bool batchnorm = false, bool additional_bias = false, bool final = false, bool deConv = false, int groups = 1); 
     virtual ~LayerWgs();
 
     int inputs, outputs;
@@ -165,7 +165,7 @@ class Conv2d : public LayerWgs {
 public:
     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, bool final = false);
+                std::string fname_weights, bool batchnorm = false, bool deConv = false, bool final = false, int groups = 1);
     virtual ~Conv2d();
     virtual layerType_t getLayerType() { return LAYER_CONV2D; };
 
@@ -173,6 +173,7 @@ public:
 
     int kernelH, kernelW, strideH, strideW, paddingH, paddingW;
     bool deConv;
+    int groups;
 
 protected:
     cudnnFilterDescriptor_t filterDesc;
@@ -196,8 +197,8 @@ class DeConv2d : public Conv2d {
 public:
     DeConv2d( Network *net, int out_ch, int kernelH, int kernelW,
             int strideH, int strideW, int paddingH, int paddingW,
-            std::string fname_weights, bool batchnorm = false) :
+            std::string fname_weights, bool batchnorm = false, int groups = 1) :
-            Conv2d(net, out_ch, kernelH, kernelW, strideH, strideW, paddingH, paddingW, fname_weights, batchnorm, true) {}
+            Conv2d(net, out_ch, kernelH, kernelW, strideH, strideW, paddingH, paddingW, fname_weights, batchnorm, true, false, groups) {}
     virtual ~DeConv2d() {}
     virtual layerType_t getLayerType() { return LAYER_DECONV2D; };
 

src/Conv2d.cpp

@@ -17,8 +17,6 @@ void Conv2d::initCUDNN(bool back) {
         idim = output_dim;
         odim = input_dim;
     }
-    //idim.print();
-    //odim.print();
 
     checkCUDNN( cudnnCreateFilterDescriptor(&filterDesc) );
     checkCUDNN( cudnnCreateConvolutionDescriptor(&convDesc) );
@@ -29,7 +27,7 @@ void Conv2d::initCUDNN(bool back) {
                                            net->tensorFormat, net->dataType, idim.n, idim.c, idim.h, idim.w) );
 
     checkCUDNN( cudnnSetFilter4dDescriptor(filterDesc,
-                                           net->dataType, net->tensorFormat, odim.c, idim.c,
+                                           net->dataType, net->tensorFormat, odim.c, idim.c/groups,
                                            kernelH, kernelW) );
 
     checkCUDNN( cudnnSetConvolution2dDescriptor(convDesc,
@@ -38,16 +36,20 @@ void Conv2d::initCUDNN(bool back) {
                                                 1,1, // upscale
                                                 CUDNN_CROSS_CORRELATION, CUDNN_DATA_FLOAT) );
 
+    checkCUDNN(  cudnnSetConvolutionGroupCount(convDesc,
+                                                groups) );
+
     // check dimension of convolution output
     dataDim_t tmpdim;
     checkCUDNN( cudnnGetConvolution2dForwardOutputDim(
             convDesc, srcTensor, filterDesc,
             &tmpdim.n, &tmpdim.c, &tmpdim.h, &tmpdim.w) );
+
     if(odim.n != tmpdim.n || odim.c != tmpdim.c || odim.h != tmpdim.h || odim.w != tmpdim.w) {
         std::cout<<"tkdim input: "; idim.print();
         std::cout<<"tkdim output: "; odim.print();
         std::cout<<"cudnndim: "; tmpdim.print();
-        FatalError("Eror conv dimension mismatch");
+        FatalError("Error conv dimension mismatch");
     }
 
     checkCUDNN( cudnnSetTensor4dDescriptor(dstTensor,
@@ -119,11 +121,10 @@ void Conv2d::inferCUDNN(dnnType* srcData, bool back) {
 
 Conv2d::Conv2d( Network *net, int out_ch, int kernelH, int kernelW,
                 int strideH, int strideW, int paddingH, int paddingW,
-                std::string fname_weights, bool batchnorm, bool deConv, bool final) :
+                std::string fname_weights, bool batchnorm, bool deConv, bool final, int groups) :
     
     LayerWgs(net, net->getOutputDim().c, out_ch, kernelH, kernelW, 1, 
-             fname_weights, batchnorm, false, final) {
+             fname_weights, batchnorm, false, final, deConv, groups) {
-
     this->kernelH = kernelH;
     this->kernelW = kernelW;
     this->strideH = strideH;
@@ -131,6 +132,7 @@ Conv2d::Conv2d( Network *net, int out_ch, int kernelH, int kernelW,
     this->paddingH = paddingH;
     this->paddingW = paddingW;
     this->deConv = deConv;
+    this->groups = groups;
 
     if(!deConv) {
         output_dim.n = input_dim.n;

src/DeformConv2d.cpp

@@ -76,8 +76,8 @@ DeformConv2d::~DeformConv2d() {
 
     checkCUDNN( cudnnDestroyTensorDescriptor(biasTensorDesc) );
     checkCuda( cudaFree(dstData) );
-    checkCuda( cudaFreeHost(ones_d1) );
+    checkCuda( cudaFree(ones_d1) );
-    checkCuda( cudaFreeHost(ones_d2) );
+    checkCuda( cudaFree(ones_d2) );
     checkCuda( cudaFree(offset) );
     checkCuda( cudaFree(mask) );
     checkCuda( cudaFree(output_conv) );

src/LayerWgs.cpp

@@ -8,8 +8,13 @@ namespace tk { namespace dnn {
 
 LayerWgs::LayerWgs(Network *net, int inputs, int outputs, 
                    int kh, int kw, int kl, 
-                   std::string fname_weights, bool batchnorm, bool additional_bias, bool final) : Layer(net, final) {
+                   std::string fname_weights, bool batchnorm, bool additional_bias, bool final, bool deConv, int groups) : Layer(net, final) {
-
+                       
+    if(deConv)
+        inputs = inputs/groups;
+    else
+        outputs = outputs/groups;    
+        
     this->inputs  = inputs;
     this->outputs = outputs;    
     this->weights_path  = std::string(fname_weights);

src/NetworkRT.cpp

@@ -245,6 +245,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
         checkNULL(lRTconv);
         lRTconv->setStride(DimsHW{l->strideH, l->strideW});
         lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
+        lRTconv->setNbGroups(l->groups);
         lRT = (ILayer*) lRTconv;
     } else {
         IDeconvolutionLayer *lRTconv = networkRT->addDeconvolution(*input, 
@@ -252,6 +253,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
         checkNULL(lRTconv);
         lRTconv->setStride(DimsHW{l->strideH, l->strideW});
         lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
+        lRTconv->setNbGroups(l->groups);
         lRT = (ILayer*) lRTconv;
         
         Dims d = lRTconv->getOutput(0)->getDimensions();