Add batchnorm

Author: pxl-th

ext/AMDGPUExt/AMDGPUExt.jl

@@ -1,6 +1,7 @@
 module AMDGPUExt
 
 import ChainRulesCore
+import ChainRulesCore: NoTangent
 import Flux
 import Flux: FluxCPUAdaptor, _amd, _isleaf, adapt_storage, fmap
 
@@ -9,6 +10,7 @@ using Adapt
 using Random
 using Zygote
 
+const MIOPENFloat = AMDGPU.MIOpen.MIOPENFloat
 const USE_AMDGPU = Ref{Union{Nothing, Bool}}(nothing)
 
 function check_use_amdgpu()

ext/AMDGPUExt/batchnorm.jl

@@ -0,0 +1,20 @@
+function (b::Flux.BatchNorm)(x::ROCArray{T}) where T <: MIOPENFloat
+    bλ.(_amd_batchnorm(x, b.γ, b.β; μ=b.μ, σ²=b.σ², ϵ=b.ϵ))
+end
+
+function _amd_batchnorm(x, γ, β; μ, σ², ϵ)
+    if NNlib.within_gradient(x)
+        return AMDGPU.MIOpen.batchnorm_training(x, γ, β, μ, σ²; ϵ, iteration=0) # TODO iteration
+    else
+        return AMDGPU.MIOpen.batchnorm_inference(x, γ, β, μ, σ²; ϵ)
+    end
+end
+
+function ChainRulesCore.rrule(::typeof(_amd_batchnorm), x, γ, β; μ, σ², ϵ)
+    y, μ_saved, ν_saved = _amd_batchnorm(x, γ, β; μ, σ², ϵ)
+    function _batchnorm_pullback(Δ)
+        dx, dγ, dβ = MIOpen.∇batchnorm(Δ, x, γ, β, μ_saved, ν_saved)
+        (NoTangent(), dx, dγ, dβ)
+    end
+    y, _batchnorm_pullback
+end

test/amd/basic.jl

@@ -78,6 +78,12 @@ end
     @test Flux.onecold(y, l) == ['a', 'a', 'a']
 end
 
+@testset "Batchnorm" begin
+    bn = BatchNorm(3, σ)
+    x = rand(Float32, 16, 16, 3, 4)
+    amdgputest(bn, x; atol=1f-3)
+end
+
 # FIXME scalar indexing. Needs NNlib.scatter?
 # @testset "Flux.onehot gpu" begin
 #     y = Flux.onehotbatch(ones(3), 1:2) |> Flux.amd