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

Change name to σ² for better consistency

Browse Source
This commit is contained in:
Elliot Saba 2019-01-20 23:57:19 +00:00
parent c5d5a5c2a8
commit a7143553df
1 changed files with 12 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
src/layers/normalise.jl
View File

@ -72,7 +72,7 @@ function Base.show(io::IO, l::LayerNorm)
end end
""" """
BatchNorm(channels::Integer, σ = identity; BatchNorm(channels::Integer, σ² = identity;
initβ = zeros, initγ = ones, initβ = zeros, initγ = ones,
ϵ = 1e-8, momentum = .1) ϵ = 1e-8, momentum = .1)
@ -102,11 +102,11 @@ m = Chain(
``` ```
""" """
mutable struct BatchNorm{F,V,W,N} mutable struct BatchNorm{F,V,W,N}
λ::F # activation function λ::F # activation function
β::V # bias β::V # bias
γ::V # scale γ::V # scale
μ::W # moving mean μ::W # moving mean
σ::W # moving std σ²::W # moving std
ϵ::N ϵ::N
momentum::N momentum::N
active::Bool active::Bool
@ -132,31 +132,31 @@ function (BN::BatchNorm)(x)
if !BN.active if !BN.active
μ = reshape(BN.μ, affine_shape...) μ = reshape(BN.μ, affine_shape...)
σ = reshape(BN.σ, affine_shape...) σ² = reshape(BN.σ², affine_shape...)
else else
T = eltype(data(x)) T = eltype(data(x))
axes = [1:dims-2; dims] # axes to reduce along (all but channels axis) axes = [1:dims-2; dims] # axes to reduce along (all but channels axis)
μ = mean(x, dims = axes) μ = mean(x, dims = axes)
meansub = (x .- μ) meansub = (x .- μ)
σ = mean(meansub .* meansub, dims = axes) σ² = mean(meansub .* meansub, dims = axes)
# update moving mean/std # update moving mean/std
mtm = convert(T, data(BN.momentum)) mtm = convert(T, data(BN.momentum))
BN.μ = (1 - mtm) .* BN.μ .+ mtm .* data(reshape(μ, :)) BN.μ = (1 - mtm) .* BN.μ .+ mtm .* data(reshape(μ, :))
BN.σ = ((1 - mtm) .* BN.σ .+ mtm .* data(reshape(σ, :)) .* m ./ (m - 1)) BN.σ² = ((1 - mtm) .* BN.σ² .+ mtm .* data(reshape(σ², :)) .* m ./ (m - 1))
end end
let λ = BN.λ, ϵ = eltype(data(σ²))(BN.ϵ) let λ = BN.λ, ϵ = eltype(data(σ²))(BN.ϵ)
λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ sqrt.(σ .+ ϵ)) .+ reshape(β, affine_shape...)) λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ sqrt.(σ² .+ ϵ)) .+ reshape(β, affine_shape...))
end end
end end
children(BN::BatchNorm) = children(BN::BatchNorm) =
(BN.λ, BN.β, BN.γ, BN.μ, BN.σ, BN.ϵ, BN.momentum, BN.active) (BN.λ, BN.β, BN.γ, BN.μ, BN.σ², BN.ϵ, BN.momentum, BN.active)
mapchildren(f, BN::BatchNorm) = # e.g. mapchildren(cu, BN) mapchildren(f, BN::BatchNorm) = # e.g. mapchildren(cu, BN)
BatchNorm(BN.λ, f(BN.β), f(BN.γ), f(BN.μ), f(BN.σ), BN.ϵ, BN.momentum, BN.active) BatchNorm(BN.λ, f(BN.β), f(BN.γ), f(BN.μ), f(BN.σ²), BN.ϵ, BN.momentum, BN.active)
_testmode!(BN::BatchNorm, test) = (BN.active = !test) _testmode!(BN::BatchNorm, test) = (BN.active = !test)