Batchnorm on GPU for Float64 values

[casper2002casper]

Some layers of GraphNeuralNetworks.jl output Float64 CuArrays, however batchnorm isn't able to work with them. Are Float64 CuArrays unsupported, or is batchnorm broken?

using Flux

println("CPU")
nn = Chain(x -> Float64.(x), BatchNorm(2)) 
data = rand(Float64, 2, 10) 
nn(data)
println("GPU")
nn = nn |> gpu
data = data |> gpu
nn(data)

CPU
GPU
ERROR: LoadError: MethodError: no method matching batchnorm(::CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ::CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ::CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ::CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ::CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ::Float32; cache=nothing, alpha=1, beta=0, eps=1.0f-5, training=false)
Closest candidates are:
  batchnorm(::CUDA.CuArray{T}, ::CUDA.CuArray{T}, ::CUDA.CuArray{T, 2}, ::CUDA.CuArray{T}, ::CUDA.CuArray{T}, ::Any; cache, alpha, beta, eps, training) where T<:Union{Float32, Float64} at ~/.julia/packages/NNlibCUDA/IeeBk/src/cudnn/batchnorm.jl:23
  batchnorm(::CUDA.CuArray{T}, ::CUDA.CuArray{T}, ::Union{CUDA.CuArray{T, 4}, CUDA.CuArray{T, 5}}, ::CUDA.CuArray{T}, ::CUDA.CuArray{T}, ::Any; cache, alpha, beta, eps, training) where T<:Union{Float32, Float64} at ~/.julia/packages/NNlibCUDA/IeeBk/src/cudnn/batchnorm.jl:27
Stacktrace:
 [1] (::BatchNorm{typeof(identity), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}})(x::CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, cache::Nothing)
   @ Flux.CUDAint ~/.julia/packages/Flux/BPPNj/src/cuda/cudnn.jl:9
 [2] (::BatchNorm{typeof(identity), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}})(x::CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer})
   @ Flux.CUDAint ~/.julia/packages/Flux/BPPNj/src/cuda/cudnn.jl:6
 [3] applychain(fs::Tuple{BatchNorm{typeof(identity), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}, x::CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}) (repeats 2 times)
   @ Flux ~/.julia/packages/Flux/BPPNj/src/layers/basic.jl:47
 [4] (::Chain{Tuple{var"#37#38", BatchNorm{typeof(identity), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}})(x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
   @ Flux ~/.julia/packages/Flux/BPPNj/src/layers/basic.jl:49
 [5] top-level scope
   @ ~/Documents/testJulia/normtest.jl:10

[terasakisatoshi]

BatchNorm(2) constructor will create an instance of struct BatchNorm its fields are type of float32. See:

Flux.jl/src/layers/normalise.jl

Lines 284 to 298 in 5e05de7

	function BatchNorm(chs::Int, λ=identity;
	initβ=zeros32, initγ=ones32,
	affine=true, track_stats=true,
	ϵ=1f-5, momentum=0.1f0)
	β = affine ? initβ(chs) : nothing
	γ = affine ? initγ(chs) : nothing
	μ = track_stats ? zeros32(chs) : nothing
	σ² = track_stats ? ones32(chs) : nothing
	return BatchNorm(λ, β, γ,
	μ, σ², ϵ, momentum,
	affine, track_stats,
	nothing, chs)
	end

You may want to pass arguments β, γ etc... directly to the default constructor of BatchNorm if you really need Float64

[DhairyaLGandhi]

Or use Flux.f64(BatchNorm(2))