Merge pull request #36 from FluxML/cl/bn

track_stats option for batchnorm

Author: CarloLucibello

ext/NNlibCUDA/src/cudnn/batchnorm.jl

@@ -13,44 +13,65 @@ end
 
 BNCache() = BNCache(nothing, nothing)
 
-@inline _wsize(y) = (fill(1, ndims(y)-2)..., size(y)[end-1], 1)
+@inline _wsize(x::AbstractArray{<:Any,N}) where N = ntuple(i -> i == N-1 ? size(x, N-1) : 1, N)
+
+function batchnorm(g::Nothing, b::Nothing, x::DenseCuArray,
+                   running_mean, running_var, momentum; kws...)
+  affine_sz = _wsize(x)
+  g = fill!(similar(x, affine_sz), 1)
+  b = fill!(similar(x, affine_sz), 0)
+  return batchnorm(g, b, x, running_mean, running_var, momentum; kws...)
+end
 
 # NOTE: CuDNN supports only 4D and 5D Tensors for BatchNorm Operations
 # so reshape a 2D Tensor into 4D
-batchnorm(g::DenseCuArray{T}, b::DenseCuArray{T}, x::DenseCuArray{T,2},
-          running_mean::DenseCuArray{T}, running_var::DenseCuArray{T}, momentum;
-          cache = nothing, alpha = T(1), beta = T(0),
-          eps = T(1e-5), training = true) where T<:Union{Float32, Float64} =
-  dropdims(batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), running_mean, running_var, momentum,
-            cache = cache, alpha = alpha, beta = beta, eps = eps, training = training), dims = (1, 2))
+function batchnorm(g::DenseCuArray{T}, b::DenseCuArray{T}, x::DenseCuArray{T,2},
+                   running_mean, running_var, momentum; kws...) where T<:CUDNNFloat
+  x = reshape(x, 1, 1, size(x, 1), size(x, 2))
+  y = batchnorm(g, b, x, running_mean, running_var, momentum; kws...)
+  return dropdims(y, dims = (1, 2))
+end
 
 function batchnorm(g::DenseCuArray{T}, b::DenseCuArray{T}, x::Union{DenseCuArray{T,4},DenseCuArray{T,5}},
-                    running_mean::DenseCuArray{T}, running_var::DenseCuArray{T}, momentum;
-                    cache = nothing, alpha = T(1), beta = T(0),
-                    eps = T(1e-5), training = true) where T<:Union{Float32, Float64}
-  cudnnBNForward!(similar(x), g, b, x, running_mean, running_var, momentum, cache = cache,
-      alpha = alpha, beta = beta, eps = eps, training = training)
+                   running_mean, running_var, momentum; kws...) where T<:CUDNNFloat
+  cudnnBNForward!(similar(x), g, b, x, running_mean, running_var, momentum; kws...)
 end
 
 function cudnnBNForward!(y::DenseCuArray{T}, g::DenseCuArray{T}, b::DenseCuArray{T}, x::DenseCuArray{T},
-                        running_mean::DenseCuArray{T}, running_var::DenseCuArray{T},
-                        momentum; cache = nothing,
+                        running_mean, running_var, momentum;
+                        cache = nothing,
                         alpha = T(1), beta = T(0),
-                        eps = T(1e-5), training = true) where T<:Union{Float32, Float64}
+                        eps = T(1e-5),
+                        training = true,
+                        affine = true,
+                        track_stats = true) where T<:CUDNNFloat
   dims = _wsize(x)
   if eps < CUDNN_BN_MIN_EPSILON
-    # warn("eps ",eps," is too small for CuDNN so eps has been assigned the value ", CUDNN_BN_MIN_EPSILON)
+    @warn "eps $eps is too small for CuDNN, setting to CUDNN_BN_MIN_EPSILON=$CUDNN_BN_MIN_EPSILON"
     eps = CUDNN_BN_MIN_EPSILON
   end
+
+  if running_mean === nothing || running_var === nothing
+    running_mean !== running_var && throw(ArgumentError("both or neither of running_mean and running_var must be nothing"))
+    if track_stats || !training
+      running_mean = fill!(similar(x, dims), 0)
+      running_var = fill!(similar(x, dims), 1)
+    end
+  end
+
   xd = cudnnTensorDescriptor(x)
   yd = cudnnTensorDescriptor(y)
   gd = cudnnTensorDescriptor(CUDNN_TENSOR_NCHW, cudnnDataType(T), Cint(length(dims)), dim4(dims,Val(CUDNN_TENSOR_NCHW)))
 
   if training
+    if !track_stats
+      running_mean = CU_NULL
+      running_var = CU_NULL
+    end
 
     if cache !== nothing
-      mean = zeros(CuArray{T}, dims...)
-      ivar = ones(CuArray{T}, dims...)
+      mean = fill!(similar(x, dims), 0)
+      ivar = fill!(similar(x, dims), 1)
     else
       mean = CU_NULL
       ivar = CU_NULL
@@ -68,57 +89,67 @@ function cudnnBNForward!(y::DenseCuArray{T}, g::DenseCuArray{T}, b::DenseCuArray
   return y
 end
 
+function ∇batchnorm(g::Nothing, b::Nothing, x::DenseCuArray, dy::DenseCuArray,
+                    running_mean, running_var, momentum; kws...)
+  affine_sz = _wsize(x)
+  g = fill!(similar(x, affine_sz), 1)
+  b = fill!(similar(x, affine_sz), 0)
+  return ∇batchnorm(g, b, x, dy, running_mean, running_var, momentum; kws...)
+end
+
 function ∇batchnorm(g::DenseCuArray{T}, b::DenseCuArray{T}, x::DenseCuArray{T, 2}, dy::DenseCuArray{T, 2},
-            running_mean::DenseCuArray{T}, running_var::DenseCuArray{T}, momentum;
-            cache = nothing, eps = T(1e-5), alpha = T(1),
-            beta = T(0), training = true) where T<:Union{Float32, Float64}
+            running_mean, running_var, momentum;
+            kws...) where T<:CUDNNFloat
   dg, db, dx = ∇batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), reshape(dy, 1, 1, size(dy, 1),
-                          size(dy, 2)), running_mean, running_var, momentum, cache = cache, eps = eps,
-                          alpha = alpha, beta = beta, training = training)
+                          size(dy, 2)), running_mean, running_var, momentum; kws...)
   (dg, db, dropdims(dx, dims = (1, 2)))
 end
 
+
 function ∇batchnorm(g::DenseCuArray{T}, b::DenseCuArray{T}, x::DenseCuArray{T}, dy::DenseCuArray{T},
-                    running_mean::DenseCuArray{T}, running_var::DenseCuArray{T}, momentum;
-                    cache = nothing, eps = T(1e-5), alpha = T(1),
-                    beta = T(0), training = true) where T<:Union{Float32, Float64}
+                    running_mean, running_var, momentum;
+                    affine=true, kws...) where T<:CUDNNFloat
   dg = similar(g)
   db = similar(b)
   dx = similar(x)
-  cudnnBNBackward!(dg, g, db, dx, x, dy, running_mean, running_var, T(momentum),
-    training = training, cache = cache, eps = eps, alpha = alpha, beta = beta)
-  (dg, db, dx)
+  cudnnBNBackward!(dg, g, db, dx, x, dy, running_mean, running_var, T(momentum); kws...)
+  if affine
+    (dg, db, dx)
+  else
+    # CUDNN always calculates dg and db, therefore we just have to drop them
+    (nothing, nothing, dx)
+  end
 end
 
 function cudnnBNBackward!(dg::DenseCuArray{T}, g::DenseCuArray{T}, db::DenseCuArray{T},
                           dx::DenseCuArray{T}, x::DenseCuArray{T}, dy::DenseCuArray{T},
-                          running_mean::DenseCuArray{T}, running_var::DenseCuArray{T},
+                          running_mean, running_var,
                           momentum; cache = nothing, eps = T(1e-5),
                           alpha = T(1), beta = T(0),
-                          dalpha = T(1), dbeta = T(0), training = true) where T<:Union{Float32, Float64}
-  if training
-    xd = cudnnTensorDescriptor(x)
-    dyd = cudnnTensorDescriptor(dy)
-    dxd = cudnnTensorDescriptor(dx)
-    gd = cudnnTensorDescriptor(CUDNN_TENSOR_NCHW, cudnnDataType(T), Cint(length(_wsize(x))), dim4(_wsize(x),Val(CUDNN_TENSOR_NCHW)))
-    if cache !== nothing
-      mean, ivar = cache.mean, cache.ivar
-      info("mean and ivar are fetched from the cache")
-    else
-      mean, ivar = CU_NULL, CU_NULL
-    end
-
-    if eps < CUDNN_BN_MIN_EPSILON
-      eps = CUDNN_BN_MIN_EPSILON
-    end
+                          dalpha = T(1), dbeta = T(0), training = true,
+                          track_stats = true) where T<:CUDNNFloat
+  if !track_stats
+    running_mean = CU_NULL
+    running_var = CU_NULL
+  end
 
-    cudnnBatchNormalizationBackward(handle(), CUDNN_BATCHNORM_SPATIAL, scalingParameter(T, alpha), scalingParameter(T, beta), scalingParameter(T, dalpha), scalingParameter(T, dbeta), xd, x, dyd, dy, dxd, dx, gd, g, dg, db, eps, mean, ivar)
+  xd = cudnnTensorDescriptor(x)
+  dyd = cudnnTensorDescriptor(dy)
+  dxd = cudnnTensorDescriptor(dx)
+  gd = cudnnTensorDescriptor(CUDNN_TENSOR_NCHW, cudnnDataType(T), Cint(length(_wsize(x))), dim4(_wsize(x),Val(CUDNN_TENSOR_NCHW)))
+  if cache !== nothing
+    @debug "fetching mean and ivar from the cache"
+    mean, ivar = cache.mean, cache.ivar
   else
-    ivar = 1 ./ sqrt.(reshape(running_var, _wsize(x)) .+ eps)
-    dx .= dy .* reshape(g, _wsize(x)) .* ivar
-    rdims = ((1:ndims(x)-2)..., ndims(x))
-    dg .= vec(sum(dy .* (x .- reshape(running_mean, _wsize(x))) .* ivar, dims=rdims))
-    db .= vec(sum(dy, dims=rdims))
+    mean, ivar = CU_NULL, CU_NULL
+  end
+
+  if eps < CUDNN_BN_MIN_EPSILON
+    @warn "eps $eps is too small for CuDNN, setting to CUDNN_BN_MIN_EPSILON=$CUDNN_BN_MIN_EPSILON"
+    eps = CUDNN_BN_MIN_EPSILON
   end
+
+  cudnnBatchNormalizationBackward(handle(), CUDNN_BATCHNORM_SPATIAL,
+        scalingParameter(T, alpha), scalingParameter(T, beta), scalingParameter(T, dalpha), scalingParameter(T, dbeta),
+        xd, x, dyd, dy, dxd, dx, gd, g, dg, db, eps, mean, ivar)
 end
-  

ext/NNlibCUDA/test/batchnorm.jl

@@ -1,8 +1,27 @@
 @testset "Batchnorm" begin
     v = CUDA.rand(Float32, 2)
     m = CUDA.rand(Float32, 2, 5)
-    for training in (false, true)
-        NNlibCUDA.batchnorm(v, v, m, v, v, 1.0; training=training)
-        NNlibCUDA.∇batchnorm(v, v, m, m, v, v, 1.0; training=training)
-    end
+
+    @testset for training in (true, false), track_stats in (true, false)
+        kws = (training=training, track_stats=track_stats)
+
+        # Normal
+        NNlibCUDA.batchnorm(v, v, m, v, v, 1.0; kws...)
+        NNlibCUDA.∇batchnorm(v, v, m, m, v, v, 1.0; kws...)
+
+        # No affine
+        NNlibCUDA.batchnorm(nothing, nothing, m, v, v, 1.0; kws...)
+        NNlibCUDA.∇batchnorm(nothing, nothing, m, m, v, v, 1.0; kws...)
+
+        # No tracking
+        NNlibCUDA.batchnorm(v, v, m, nothing, nothing, 1.0; kws...)
+        NNlibCUDA.∇batchnorm(v, v, m, m, nothing, nothing, 1.0; kws...)
+
+        # Both or neither tracked or affine params must be set
+        for (α, β) in ((v, nothing), (nothing, v))
+            @test_throws MethodError NNlibCUDA.batchnorm(α, β, m, v, v, 1.0; kws...)
+            @test_throws MethodError NNlibCUDA.∇batchnorm(α, β, m, m, v, v, 1.0; kws...)
+            @test_throws ArgumentError NNlibCUDA.batchnorm(v, v, m, α, β, 1.0; kws...)
+        end
+    end 
 end