From 33f36ab2048b7746a018fbe97d95a59d2e515fd9 Mon Sep 17 00:00:00 2001
From: fbagni <gattinomicino>
Date: Wed, 30 Oct 2019 09:40:29 +0100
Subject: [PATCH] Deconv tensorrt

---
 include/tkDNN/Layer.h     |  3 +--
 include/tkDNN/NetworkRT.h |  1 +
 src/NetworkRT.cpp         | 25 +++++++++++++++++--------
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/include/tkDNN/Layer.h b/include/tkDNN/Layer.h
index 483b7b5..6015235 100644
--- a/include/tkDNN/Layer.h
+++ b/include/tkDNN/Layer.h
@@ -161,6 +161,7 @@ public:
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
 
     int kernelH, kernelW, strideH, strideW, paddingH, paddingW;
+    bool deConv;
 
 protected:
     cudnnFilterDescriptor_t filterDesc;
@@ -173,8 +174,6 @@ protected:
     void inferCUDNN(dnnType* srcData, bool back = false);
     void*  workSpace;
     size_t ws_sizeInBytes;
-
-    bool deConv;
 };
 
 
diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h
index d30da03..97dd478 100644
--- a/include/tkDNN/NetworkRT.h
+++ b/include/tkDNN/NetworkRT.h
@@ -75,6 +75,7 @@ public:
 
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Layer *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Conv2d *l);
+    nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, DeConv2d *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Activation *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Dense *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Pooling *l);
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index 2b3319c..ce29f01 100644
--- a/src/NetworkRT.cpp
+++ b/src/NetworkRT.cpp
@@ -159,7 +159,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
 
     if(type == LAYER_DENSE)
         return convert_layer(input, (Dense*) l);
-    if(type == LAYER_CONV2D)
+    if(type == LAYER_CONV2D || type == LAYER_DECONV2D)
         return convert_layer(input, (Conv2d*) l);
     if(type == LAYER_POOLING)
         return convert_layer(input, (Pooling*) l);
@@ -232,13 +232,22 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
     else
         b = { dtRT, nullptr, 0}; //on batchnorm bias are added later
 
-    // Add a convolution layer with 20 outputs and a 5x5 filter.
-    IConvolutionLayer *lRT = networkRT->addConvolution(*input, 
-               l->outputs, DimsHW{l->kernelH, l->kernelW}, w, b);
-    checkNULL(lRT);
-
-    lRT->setStride(DimsHW{l->strideH, l->strideW});
-    lRT->setPadding(DimsHW{l->paddingH, l->paddingW});
+    ILayer *lRT = nullptr;
+    if(!l->deConv) {
+        IConvolutionLayer *lRTconv = networkRT->addConvolution(*input, 
+            l->outputs, DimsHW{l->kernelH, l->kernelW}, w, b);
+        checkNULL(lRT);
+        lRTconv->setStride(DimsHW{l->strideH, l->strideW});
+        lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
+        lRT = (ILayer*) lRTconv;
+    } else {
+        IDeconvolutionLayer *lRTconv = networkRT->addDeconvolution(*input, 
+            l->outputs, DimsHW{l->kernelH, l->kernelW}, w, b);
+        checkNULL(lRT);
+        lRTconv->setStride(DimsHW{l->strideH, l->strideW});
+        lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
+        lRT = (ILayer*) lRTconv;
+    }
 
     if(l->batchnorm) {
         Weights power{dtRT, power_b, l->outputs};