TastyPancakes/Flux.jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

batchnorm: parameterize momentum and epsilon

Browse Source
This commit is contained in:
Iblis Lin 2017-10-30 11:34:51 +08:00
parent 669273b008
commit e0201be770
1 changed files with 11 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
src/layers/normalisation.jl
View File

@ -67,32 +67,35 @@ julia> m = Chain(
Chain(Dense(784, 64), BatchNorm(64, λ = NNlib.relu), Dense(64, 10), BatchNorm(10), NNlib.softmax) Chain(Dense(784, 64), BatchNorm(64, λ = NNlib.relu), Dense(64, 10), BatchNorm(10), NNlib.softmax)
``` ```
""" """
mutable struct BatchNorm{F,V} mutable struct BatchNorm{F,V,N}
λ::F # activation function λ::F # activation function
β::V # bias β::V # bias
γ::V # scale γ::V # scale
μ # moving mean μ # moving mean
σ # moving std σ # moving std
ϵ::Float64 ϵ::N
momentum::Float64 momentum::N
active::Bool active::Bool
end end
BatchNorm(dims::Integer...; λ = identity, BatchNorm(dims::Integer...; λ = identity,
initβ = zeros, initγ = ones, ϵ = 1e-8, momentum = .1) = initβ = zeros, initγ = ones, ϵ = 1e-8, momentum = .1) =
BatchNorm(λ, param(initβ(dims)), param(initγ(dims)), 0., 1., momentum, ϵ, true) BatchNorm(λ, param(initβ(dims)), param(initγ(dims)), 0., 1., ϵ, momentum, true)
function (BN::BatchNorm)(x) function (BN::BatchNorm)(x)
if !BN.active if !BN.active
μ = BN.μ μ = BN.μ
σ = BN.σ σ = BN.σ
else else
T = eltype(x)
ϵ = T(BN.ϵ)
m = size(x, 2) # batch size m = size(x, 2) # batch size
μ = sum(x, 2) ./ m μ = sum(x, 2) ./ m
σ = sqrt.(sum((x .- μ).^2, 2) ./ (m - 1) .+ BN.ϵ) σ = sqrt.(sum((x .- μ).^2, 2) ./ (m - 1) .+ ϵ)
# update moving mean/std # update moving mean/std
mtm = BN.momentum mtm = T(BN.momentum)
BN.μ = mtm .* μ.data .+ (1 - mtm) .* BN.μ BN.μ = mtm .* μ.data .+ (1 - mtm) .* BN.μ
BN.σ = mtm .* σ.data .+ (1 - mtm) .* BN.σ BN.σ = mtm .* σ.data .+ (1 - mtm) .* BN.σ
end end
@ -102,7 +105,8 @@ end
children(BN::BatchNorm) = children(BN::BatchNorm) =
(BN.λ, BN.β, BN.γ, BN.μ, BN.σ, BN.momentum, BN.ϵ, BN.active) (BN.λ, BN.β, BN.γ, BN.μ, BN.σ, BN.momentum, BN.ϵ, BN.active)
mapchildren(f, BN::BatchNorm) =
mapchildren(f, BN::BatchNorm) = # e.g. mapchildren(cu, BN)
BatchNorm(λ, f(BN.β), f(BN.γ), BN.μ, BN.σ, BN.momentum, BN.ϵ, BN.active) BatchNorm(λ, f(BN.β), f(BN.γ), BN.μ, BN.σ, BN.momentum, BN.ϵ, BN.active)
_testmode!(BN::BatchNorm, test) = (BN.active = !test) _testmode!(BN::BatchNorm, test) = (BN.active = !test)