From c5d5a5c2a8fe0e20ddeb6a01aab497e7480052d2 Mon Sep 17 00:00:00 2001
From: Elliot Saba <staticfloat@gmail.com>
Date: Fri, 18 Jan 2019 15:12:49 -0500
Subject: [PATCH] Cleanup BatchNorm implementation

This provides greater datatype persistence
---
 src/layers/normalise.jl | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/layers/normalise.jl b/src/layers/normalise.jl
index 85224ed2..c25fe798 100644
--- a/src/layers/normalise.jl
+++ b/src/layers/normalise.jl
@@ -134,21 +134,21 @@ function (BN::BatchNorm)(x)
     μ = reshape(BN.μ, affine_shape...)
     σ = reshape(BN.σ, affine_shape...)
   else
-    T = eltype(x)
+    T = eltype(data(x))
 
-    ϵ = data(convert(T, BN.ϵ))
     axes = [1:dims-2; dims] # axes to reduce along (all but channels axis)
     μ = mean(x, dims = axes)
-    σ = sqrt.(mean((x .- μ).^2, dims = axes) .+ ϵ)
+    meansub = (x .- μ)
+    σ = mean(meansub .* meansub, dims = axes)
 
     # update moving mean/std
-    mtm = data(convert(T, BN.momentum))
-    BN.μ = (1 - mtm) .* BN.μ .+ mtm .* dropdims(data(μ), dims = (axes...,))
-    BN.σ = (1 - mtm) .* BN.σ .+ mtm .* dropdims(data(σ), dims = (axes...,)) .* m ./ (m - 1)
+    mtm = convert(T, data(BN.momentum))
+    BN.μ = (1 - mtm) .* BN.μ .+ mtm .* data(reshape(μ, :))
+    BN.σ = ((1 - mtm) .* BN.σ .+ mtm .* data(reshape(σ, :)) .* m ./ (m - 1))
   end
 
-  let λ = BN.λ
-    λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ σ) .+ reshape(β, affine_shape...))
+  let λ = BN.λ, ϵ = eltype(data(σ²))(BN.ϵ)
+    λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ sqrt.(σ .+ ϵ)) .+ reshape(β, affine_shape...))
   end
 end