# BatchNorm: only track\_stats=true supported on gpu

**URL:** <https://discourse.julialang.org/t/batchnorm-only-track-stats-true-supported-on-gpu/62091>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [May 30, 2021, 1:51pm UTC](https://discourse.julialang.org/t/batchnorm-only-track-stats-true-supported-on-gpu/62091 "2021-05-30T13:51:54Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![ymtoo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ymtoo/32/18859_2.png) [@ymtoo](https://discourse.julialang.org/u/ymtoo)\
**Post date:** [May 30, 2021, 1:51pm UTC](https://discourse.julialang.org/t/batchnorm-only-track-stats-true-supported-on-gpu/62091/1 "2021-05-30T13:51:54Z")

</div>

Not able to use `BatchNorm` with `track_stats=false` on GPU.

```julia
julia> using Flux

julia> model = Chain(Dense(10,5),BatchNorm(5,relu; track_stats=false))
Chain(Dense(10, 5), BatchNorm(5, relu))

julia> modelgpu = gpu(model)
Chain(Dense(10, 5), BatchNorm(5, relu))

julia> x = randn(Float32, 10, 100);

julia> model(x)
5×100 Matrix{Float32}:
 0.897995 1.41563 0.0 0.625536 0.0 … 0.358508 0.302867 0.0 0.0
 0.903481 0.0 0.782808 0.673593 0.0 1.81217 0.277665 0.0 0.0
 1.37133 0.0 1.94733 0.259229 0.0 1.73941 0.0 0.137841 0.287855
 0.179537 0.0 0.316961 0.0 0.446272 0.0 0.0 0.0 0.0
 0.721404 1.4531 0.0 0.0 0.183346 0.0 0.0 1.31021 0.511247

julia> modelgpu(gpu(x))
ERROR: AssertionError: BatchNorm: only track_stats=true supported on gpu
Stacktrace:
 [1] (::BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1}, Float32, Nothing})(x::CUDA.CuArray{Float32, 2}, cache::Nothing)
   @ Flux.CUDAint ~/.julia/packages/Flux/6o4DQ/src/cuda/cudnn.jl:7
 [2] (::BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1}, Float32, Nothing})(x::CUDA.CuArray{Float32, 2})
   @ Flux.CUDAint ~/.julia/packages/Flux/6o4DQ/src/cuda/cudnn.jl:6
 [3] applychain(fs::Tuple{BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1}, Float32, Nothing}}, x::CUDA.CuArray{Float32, 2}) (repeats 2 times)
   @ Flux ~/.julia/packages/Flux/6o4DQ/src/layers/basic.jl:36
 [4] (::Chain{Tuple{Dense{typeof(identity), CUDA.CuArray{Float32, 2}, CUDA.CuArray{Float32, 1}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1}, Float32, Nothing}}})(x::CUDA.CuArray{Float32, 2})
   @ Flux ~/.julia/packages/Flux/6o4DQ/src/layers/basic.jl:38
 [5] top-level scope
   @ REPL[40]:1
 [6] top-level scope
   @ ~/.julia/packages/CUDA/3VnCC/src/initialization.jl:81

```

Julia and package versions:

```julia
julia> versioninfo()
Julia Version 1.6.1
Commit 6aaedecc44 (2021-04-23 05:59 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Core(TM) i7-10510U CPU @ 1.80GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-11.0.1 (ORCJIT, skylake)

(jl_VSWxj0) pkg> st
      Status `/tmp/jl_VSWxj0/Project.toml`
  [587475ba] Flux v0.12.3

```

Any idea why only `track_stats=true` is supported on GPU?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [May 30, 2021, 4:20pm UTC](https://discourse.julialang.org/t/batchnorm-only-track-stats-true-supported-on-gpu/62091/2 "2021-05-30T16:20:05Z")

</div>

cuDNN will unconditionally track stats for batchnorm if you invoke `cudnnBatchNormalizationForwardTraining` (which Flux does). I’m not sure if it’s still possible to have the proper behaviour if `cudnnBatchNormalizationForwardInference` is used instead when `track_stats=false`, but that would be the path forward here.

---

<div class="post-metadata">

**Author:** ![ymtoo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ymtoo/32/18859_2.png) [@ymtoo](https://discourse.julialang.org/u/ymtoo)\
**Post date:** [May 31, 2021, 3:50am UTC](https://discourse.julialang.org/t/batchnorm-only-track-stats-true-supported-on-gpu/62091/3 "2021-05-31T03:50:10Z")

</div>

Thanks for pointing out the functions of `BatchNorm` on GPU.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [May 31, 2021, 5:17am UTC](https://discourse.julialang.org/t/batchnorm-only-track-stats-true-supported-on-gpu/62091/4 "2021-05-31T05:17:30Z")

</div>

I’ve opened [Allow BatchNorm training on CUDA with `track_stats=false` · Issue #1606 · FluxML/Flux.jl · GitHub](https://github.com/FluxML/Flux.jl/issues/1606). No promises on any kind of timeline, but it’s documented now. Thanks for reporting!