From 4a1c4a5f919924f3a52b92091f720943ed8784dd Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Sun, 19 May 2024 15:13:22 +0200
Subject: [PATCH] Add epsilon param to batch norm

---
 include/layers/batch_norm.cuh  | 2 +-
 src/layers/batch_norm.cu       | 2 +-
 test/layers/test_batch_norm.cu | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/include/layers/batch_norm.cuh b/include/layers/batch_norm.cuh
index cd4b623..c24a946 100644
--- a/include/layers/batch_norm.cuh
+++ b/include/layers/batch_norm.cuh
@@ -10,7 +10,7 @@ namespace CUDANet::Layers {
 
 class BatchNorm : public WeightedLayer {
   public:
-    BatchNorm(int inputSize, int inputChannels, ActivationType activationType);
+    BatchNorm(int inputSize, int inputChannels, float epsilon, ActivationType activationType);
 
     ~BatchNorm();
 
diff --git a/src/layers/batch_norm.cu b/src/layers/batch_norm.cu
index 70e532a..486a728 100644
--- a/src/layers/batch_norm.cu
+++ b/src/layers/batch_norm.cu
@@ -12,6 +12,7 @@ using namespace CUDANet::Layers;
 BatchNorm::BatchNorm(
     int            inputSize,
     int            inputChannels,
+    float          epsilon,
     ActivationType activationType
 )
     : inputSize(inputSize), inputChannels(inputChannels) {
@@ -47,7 +48,6 @@ BatchNorm::BatchNorm(
     CUDA_CHECK(cudaMemcpy(d_length, &length, sizeof(float), cudaMemcpyHostToDevice));
 
     d_epsilon = nullptr;
-    float epsilon = 1e-5f;
     CUDA_CHECK(cudaMalloc((void **)&d_epsilon, sizeof(float)));
     CUDA_CHECK(cudaMemcpy(d_epsilon, &epsilon, sizeof(float), cudaMemcpyHostToDevice));
 
diff --git a/test/layers/test_batch_norm.cu b/test/layers/test_batch_norm.cu
index da54619..46da0fa 100644
--- a/test/layers/test_batch_norm.cu
+++ b/test/layers/test_batch_norm.cu
@@ -13,7 +13,7 @@ TEST(BatchNormLayerTest, BatchNormSmallForwardTest) {
     cudaError_t cudaStatus;
 
     CUDANet::Layers::BatchNorm batchNorm(
-        inputSize, nChannels, CUDANet::Layers::ActivationType::NONE
+        inputSize, nChannels, 1e-5f, CUDANet::Layers::ActivationType::NONE
     );
 
     std::vector<float> weights = {0.63508f, 0.64903f};