From 33844c1ab2e811cf06a0b6b26a6238c3cca6ff6d Mon Sep 17 00:00:00 2001
From: xavier <micaelaverucchi@gmail.com>
Date: Wed, 15 Jan 2020 18:06:02 +0100
Subject: [PATCH] Batchnorm eps fix, works on jetpack 4.3

---
 README.md        |  7 ++++---
 include/Layer.h  |  2 ++
 src/Conv2d.cpp   | 10 +++++-----
 src/LayerWgs.cpp |  2 +-
 src/Network.cpp  | 17 +++++++++--------
 5 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index aaf5872..cd06992 100644
--- a/README.md
+++ b/README.md
@@ -3,9 +3,10 @@ tkDNN is a Deep Neural Network library built with cuDNN primitives specifically
 The main scope is to do high performance inference on already trained models.
 
 this branch actually work on every NVIDIA GPU that support the dependencies:
-* CUDA 9
-* CUDNN 7.105
-* TENSORRT 4.02
+* CUDA 10.0
+* CUDNN 7.603
+* TENSORRT 6.01
+* OPENCV 4.1
 
 ## Dependencies
 
diff --git a/include/Layer.h b/include/Layer.h
index b882619..396288e 100644
--- a/include/Layer.h
+++ b/include/Layer.h
@@ -27,6 +27,8 @@ enum layerType_t
     LAYER_YOLO
 };
 
+#define TKDNN_BN_MIN_EPSILON 1e-5
+
 /** 
     Simple layer Father class
 */
diff --git a/src/Conv2d.cpp b/src/Conv2d.cpp
index bf13a33..9e3aeeb 100644
--- a/src/Conv2d.cpp
+++ b/src/Conv2d.cpp
@@ -120,11 +120,11 @@ dnnType *Conv2d::infer(dataDim_t &dim, dnnType *srcData)
         float one = 1;
         float zero = 0;
         cudnnBatchNormalizationForwardInference(net->cudnnHandle,
-                                                CUDNN_BATCHNORM_SPATIAL, &one, &zero,
-                                                dstTensorDesc, dstData, dstTensorDesc,
-                                                dstData, biasTensorDesc, //same tensor descriptor as bias
-                                                scales_d, bias_d, mean_d, variance_d,
-                                                CUDNN_BN_MIN_EPSILON);
+            CUDNN_BATCHNORM_SPATIAL, &one, &zero, 
+            dstTensorDesc, dstData, dstTensorDesc, 
+            dstData, biasTensorDesc, //same tensor descriptor as bias 
+            scales_d, bias_d, mean_d, variance_d, 
+            TKDNN_BN_MIN_EPSILON);
     }
     //update data dimensions
     dim = output_dim;
diff --git a/src/LayerWgs.cpp b/src/LayerWgs.cpp
index ec6cab8..869ef69 100644
--- a/src/LayerWgs.cpp
+++ b/src/LayerWgs.cpp
@@ -34,7 +34,7 @@ LayerWgs::LayerWgs(Network *net, int inputs, int outputs,
         seek += outputs;
         readBinaryFile(weights_path.c_str(), outputs, &variance_h, &variance_d, seek);
 
-        float eps = CUDNN_BN_MIN_EPSILON;
+        float eps = TKDNN_BN_MIN_EPSILON;
 
         power_h = new dnnType[outputs];
         for (int i = 0; i < outputs; i++)
diff --git a/src/Network.cpp b/src/Network.cpp
index 7ceeca5..b0ca487 100644
--- a/src/Network.cpp
+++ b/src/Network.cpp
@@ -21,10 +21,6 @@ Network::Network(dataDim_t input_dim)
               << ", CUDNN v" << cu_ver << ")\n";
     dataType = CUDNN_DATA_FLOAT;
     tensorFormat = CUDNN_TENSOR_NCHW;
-
-    checkCUDNN(cudnnCreate(&cudnnHandle));
-    checkERROR(cublasCreate(&cublasHandle));
-
     num_layers = 0;
 
     fp16 = false;
@@ -39,11 +35,16 @@ Network::Network(dataDim_t input_dim)
             fp16 = true;
         }
     }
+   
+    if(fp16)
+        std::cout<<COL_REDB<<"!! FP16 INERENCE ENABLED !!"<<COL_END<<"\n";
+    if(dla)
+        std::cout<<COL_GREENB<<"!! DLA INERENCE ENABLED !!"<<COL_END<<"\n";
+
+
+    checkCUDNN( cudnnCreate(&cudnnHandle) );
+    checkERROR( cublasCreate(&cublasHandle) );
 
-    if (fp16)
-        std::cout << COL_REDB << "!! FP16 INERENCE ENABLED !!" << COL_END << "\n";
-    if (dla)
-        std::cout << COL_GREENB << "!! DLA INERENCE ENABLED !!" << COL_END << "\n";
 }
 
 Network::~Network()