From bcf0c4eab3d63a2a01ea73b898ee032c98f400c0 Mon Sep 17 00:00:00 2001
From: perseusdg <harshvardhan.chandira@gmail.com>
Date: Mon, 10 Jan 2022 18:43:12 +0530
Subject: [PATCH] Added BatchNorm Layer (Testing still needs to be done)

---
 include/tkDNN/Layer.h | 23 ++++++++++++++-
 src/BatchNorm.cpp     | 67 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 src/BatchNorm.cpp

diff --git a/include/tkDNN/Layer.h b/include/tkDNN/Layer.h
index 28a8817..62d5a3e 100644
--- a/include/tkDNN/Layer.h
+++ b/include/tkDNN/Layer.h
@@ -32,7 +32,8 @@ enum layerType_t {
     LAYER_UPSAMPLE,
     LAYER_REGION,
     LAYER_YOLO,
-    LAYER_PADDING
+    LAYER_PADDING,
+    LAYER_BATCHNORM
 };
 
 #define TKDNN_BN_MIN_EPSILON 1e-5
@@ -89,6 +90,7 @@ public:
             case LAYER_REGION:              return "Region";
             case LAYER_YOLO:                return "Yolo";
             case LAYER_PADDING:             return "Padding";
+            case LAYER_BATCHNORM:           return "BatchNorm";
             default:                        return "unknown";
         }
     }
@@ -604,6 +606,25 @@ public:
 
 };
 
+
+class BatchNorm : public LayerBNWgs {
+public:
+    BatchNorm(Network *net,int output,std::string fname_weights);
+    virtual ~BatchNorm();
+    virtual layerType_t getLayerType(){return LAYER_BATCHNORM;};
+    virtual dnnType* infer(dataDim_t& dim,dnnType* srcData);
+    std::string weights_bin;
+protected:
+    cudnnFilterDescriptor_t filterDesc;
+    cudnnConvolutionFwdAlgoPerf_t     algo;
+    cudnnConvolutionBwdDataAlgoPerf_t bwAlgo;
+    cudnnTensorDescriptor_t biasTensorDesc;
+
+    void initCUDNN();
+    void inferCUDNN(dnnType* srcData);
+    void*  workSpace;
+    size_t ws_sizeInBytes;
+};
 /**
     Softmax layer
 */
diff --git a/src/BatchNorm.cpp b/src/BatchNorm.cpp
new file mode 100644
index 0000000..b406011
--- /dev/null
+++ b/src/BatchNorm.cpp
@@ -0,0 +1,67 @@
+#include <iostream>
+
+#include "Layer.h"
+
+namespace tk { namespace dnn { 
+    void BatchNorm::initCUDNN(){
+        cudnnTensorDescriptor_t srcTensor = srcTensorDesc;
+        cudnnTensorDescriptor_t dstTensor = dstTensorDesc;
+        dataDim_t idim,odim;
+        idim = input_dim;
+        odim = output_dim;
+
+        checkCUDNN( cudnnSetTensor4dDescriptor(srcTensor,
+                                           net->tensorFormat, net->dataType, idim.n, idim.c, idim.h, idim.w) );
+
+        checkCUDNN( cudnnCreateTensorDescriptor(&biasTensorDesc) );
+        
+        checkCUDNN( cudnnSetTensor4dDescriptor(dstTensor,
+                                           net->tensorFormat, net->dataType, odim.n, odim.c, odim.h, odim.w) );
+
+        checkCUDNN( cudnnSetTensor4dDescriptor(biasTensorDesc,
+                                           net->tensorFormat, net->dataType,
+                                           1, output_dim.c, 1, 1) );
+
+
+    }
+
+    void BatchNorm::inferCUDNN(float *srcData){
+        dnnType alpha = dnnType(1);
+        dnnType beta  = dnnType(0);
+
+        alpha = dnnType(1);
+        beta  = dnnType(1);
+        checkCUDNN( cudnnAddTensor(net->cudnnHandle,
+                                   &alpha, biasTensorDesc, bias_d,
+                                   &beta, dstTensorDesc, dstData) );
+        alpha = dnnType(1);
+        beta  = dnnType(0);
+        checkCUDNN( cudnnBatchNormalizationForwardInference(net->cudnnHandle,
+                                                    CUDNN_BATCHNORM_SPATIAL, &alpha, &beta,
+                                                    dstTensorDesc, dstData, dstTensorDesc,
+                                                    dstData, biasTensorDesc, //same tensor descriptor as bias
+                                                    scales_d, bias_d, mean_d, variance_d,
+                                                    TKDNN_BN_MIN_EPSILON) );
+    }
+
+    BatchNorm::BatchNorm(Network *net,int output,std::string fname_weights) : 
+        LayerBNWgs(net,net->getOutputDim().c,output,fname_weights){
+            output_dim = input_dim;
+            initCUDNN();
+            checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) );
+
+    }
+
+    dnnType* BatchNorm::infer(dataDim_t &dim,dnnType* srcData){
+        inferCUDNN(srcData);
+
+        dim = output_dim;
+        return dstData;
+    }
+
+    BatchNorm::~BatchNorm(){
+        checkCUDNN( cudnnDestroyTensorDescriptor(biasTensorDesc) );
+        checkCuda( cudaFree(dstData) );
+    }
+
+}}
\ No newline at end of file