Skip to content

Commit 711eb07

Browse files
pxl-thpxl-th
authored
Bring back batchnorm (#637)
1 parent 7125f2e commit 711eb07

File tree

4 files changed

+155
-0
lines changed

4 files changed

+155
-0
lines changed

‎src/dnn/MIOpen.jl

+1
Original file line numberDiff line numberDiff line change
@@ -89,5 +89,6 @@ include("descriptors.jl")
8989
include("convolution.jl")
9090
include("pooling.jl")
9191
include("activations.jl")
92+
include("batchnorm.jl")
9293

9394
end

‎src/dnn/batchnorm.jl

+124
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
"""
2+
batchnorm_training(
3+
x::ROCArray{T},
4+
γ::ROCArray{T}, β::ROCArray{T},
5+
μ::ROCArray{T}, ν::ROCArray{T}; iteration::Int, ϵ::Float64 = 1e-5,
6+
) where T <: MIOPENFloat
7+
# Arguments:
8+
- `γ`: Scaling.
9+
- `β`: Bias.
10+
- `μ`: Running mean for inference.
11+
- `ν`: Running variance for inference.
12+
If `x` has `N` dims, then `N - 1` is considered as 'feature' dimension.
13+
Meaning, γ, β, μ, ν must have `size(x, N - 1)` shape.
14+
"""
15+
function batchnorm_training(
16+
x::ROCArray{T},
17+
γ::ROCArray{T}, β::ROCArray{T},
18+
μ::ROCArray{T}, ν::ROCArray{T}; iteration::Int, ϵ::Float64 = 1e-5,
19+
) where T <: MIOPENFloat
20+
y = similar(x)
21+
22+
nd = ndims(x)
23+
n_features = size(x, nd - 1)
24+
mode = nd == 2 ? miopenBNPerActivation : miopenBNSpatial
25+
xdesc, ydesc = TensorDescriptor4D.((x, y))
26+
27+
bndesc = derive_beta_gamma_descriptors(xdesc, mode)
28+
factor = 1.0 / (1.0 + Float64(iteration))
29+
# For backward pass.
30+
μ_saved, ν_saved = similar(x, n_features), similar(x, n_features)
31+
32+
(; handle, stream) = lib_state()
33+
miopenBatchNormalizationForwardTraining(
34+
handle, mode, Ref{Float32}(1f0), Ref{Float32}(0f0),
35+
xdesc.handle, x, ydesc.handle, y, bndesc.handle, γ, β, factor,
36+
μ, ν, ϵ, μ_saved, ν_saved) |> check
37+
y, μ_saved, ν_saved
38+
end
39+
40+
"""
41+
batchnorm_inference(
42+
x::ROCArray{T},
43+
γ::ROCArray{T}, β::ROCArray{T},
44+
μ::ROCArray{T}, ν::ROCArray{T}; ϵ::Float64 = 1e-5,
45+
) where T <: MIOPENFloat
46+
# Arguments:
47+
- `γ`: Scaling.
48+
- `β`: Bias.
49+
- `μ`: Running mean for inference.
50+
- `ν`: Running variance for inference.
51+
If `x` has `N` dims, then `N - 1` is considered as 'feature' dimension.
52+
Meaning, γ, β, μ, ν must have `size(x, N - 1)` shape.
53+
"""
54+
function batchnorm_inference(
55+
x::ROCArray{T},
56+
γ::ROCArray{T}, β::ROCArray{T},
57+
μ::ROCArray{T}, ν::ROCArray{T}; ϵ::Float64 = 1e-5,
58+
) where T <: MIOPENFloat
59+
y = similar(x)
60+
61+
nd = ndims(x)
62+
mode = nd == 2 ? miopenBNPerActivation : miopenBNSpatial
63+
xdesc, ydesc = TensorDescriptor4D.((x, y))
64+
bndesc = derive_beta_gamma_descriptors(xdesc, mode)
65+
66+
(; handle, stream) = lib_state()
67+
miopenBatchNormalizationForwardInference(
68+
handle, mode, Ref{Float32}(1f0), Ref{Float32}(0f0),
69+
xdesc.handle, x, ydesc.handle, y, bndesc.handle,
70+
γ, β, μ, ν, ϵ) |> check
71+
y
72+
end
73+
74+
function ∇batchnorm(
75+
dy::ROCArray{T, N}, x::ROCArray{T, N},
76+
γ::ROCArray{T}, β::ROCArray{T},
77+
μ_saved::ROCArray{T}, ν_saved::ROCArray{T}; ϵ::Float64 = 1e-5,
78+
) where {T <: MIOPENFloat, N}
79+
dx, dγ, dβ = similar(x), similar(γ), similar(β)
80+
81+
nd = ndims(x)
82+
mode = nd == 2 ? miopenBNPerActivation : miopenBNSpatial
83+
xdesc, dxdesc, dydesc = TensorDescriptor4D.((x, dx, dy))
84+
bndesc = derive_beta_gamma_descriptors(xdesc, mode)
85+
86+
(; handle, stream) = lib_state()
87+
miopenBatchNormalizationBackward(
88+
handle, mode,
89+
Ref{Float32}(1f0), Ref{Float32}(0f0),
90+
Ref{Float32}(1f0), Ref{Float32}(0f0),
91+
xdesc.handle, x, dydesc.handle, dy, dxdesc.handle, dx,
92+
bndesc.handle, γ, dγ, dβ, ϵ, μ_saved, ν_saved) |> check
93+
dx, dγ, dβ
94+
end
95+
96+
function derive_beta_gamma_descriptors(
97+
xdesc::TensorDescriptor, mode::miopenBatchNormMode_t,
98+
)
99+
handle_ref = Ref{miopenTensorDescriptor_t}()
100+
miopenCreateTensorDescriptor(handle_ref) |> check
101+
handle = handle_ref[]
102+
103+
miopenDeriveBNTensorDescriptor(handle, xdesc.handle, mode) |> check
104+
dtype, dims, stride = unpack(handle, ndims(handle))
105+
106+
bndesc = TensorDescriptor(handle, dtype)
107+
finalizer(bndesc) do d_
108+
miopenDestroyTensorDescriptor(d_.handle) |> check
109+
end
110+
bndesc
111+
end
112+
113+
# Unsqueeze dimensions at the beginning:
114+
# _bn_expand_dims((3, 2), 4) -> (1, 1, 3, 2)
115+
function _bn_expand_dims(v, ndims)
116+
reverse(ntuple(
117+
i -> (i ≤ length(v)) ? Int64(v[end - i + 1]) : 1,
118+
Val{ndims}()))
119+
end
120+
121+
function TensorDescriptor4D(x)
122+
nd = ndims(x)
123+
TensorDescriptor(reshape(x, _bn_expand_dims(size(x), max(4, nd))))
124+
end

‎test/dnn/batchnorm.jl

+26
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
@testset "Different input dimensions" begin
2+
for sz in ((3, 2), (4, 3, 2), (5, 4, 3, 2))
3+
x = AMDGPU.ones(Float32, sz)
4+
γ = AMDGPU.ones(Float32, sz[end - 1])
5+
β = AMDGPU.zeros(Float32, sz[end - 1])
6+
μ = AMDGPU.zeros(Float32, sz[end - 1])
7+
ν = AMDGPU.ones(Float32, sz[end - 1])
8+
9+
y = MIOpen.batchnorm_inference(x, γ, β, μ, ν)
10+
@test all(isapprox.(Array(y), 1; atol=1f-5))
11+
12+
y, μ_saved, ν_saved = MIOpen.batchnorm_training(
13+
x, γ, β, μ, ν; iteration=0)
14+
@test all(isapprox.(Array(y), 0; atol=1f-5))
15+
16+
hμ_saved, hν_saved = Array(μ_saved), Array(ν_saved)
17+
18+
dy = AMDGPU.ones(Float32, size(y))
19+
dx, dγ, dβ = MIOpen.∇batchnorm(dy, x, γ, β, μ_saved, ν_saved)
20+
21+
hμ_saved2, hν_saved2 = Array(μ_saved), Array(ν_saved)
22+
23+
@test hμ_saved ≈ hμ_saved2
24+
@test hν_saved ≈ hν_saved2
25+
end
26+
end

‎test/dnn/miopen.jl

+4
Original file line numberDiff line numberDiff line change
@@ -31,4 +31,8 @@ end
3131
include("activations.jl")
3232
end
3333

34+
@testset "Batchnorm" begin
35+
include("activations.jl")
36+
end
37+
3438
end

0 commit comments

Comments
 (0)