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

Fix type instability when loading onto a gpu.

Browse Source 
Also fixes Issue #216.
This commit is contained in:
Brad Safnuk 2018-03-22 21:32:32 -04:00
parent 07b0f95d61
commit 35299d4621
1 changed files with 10 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
src/layers/normalisation.jl
View File

@ -101,12 +101,12 @@ m = Chain(
Dense(288, 10), softmax) |> gpu Dense(288, 10), softmax) |> gpu
``` ```
""" """
mutable struct BatchNorm{F,V,N} mutable struct BatchNorm{F,V, W,N}
λ::F # activation function λ::F # activation function
β::V # bias β::V # bias
γ::V # scale γ::V # scale
μ # moving mean μ::W # moving mean
σ # moving std σ::W # moving std
ϵ::N ϵ::N
momentum::N momentum::N
active::Bool active::Bool
@ -114,7 +114,7 @@ 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)), zeros(dims), zeros(dims), ϵ, momentum, true)
function (BN::BatchNorm)(x) function (BN::BatchNorm)(x)
λ, γ, β = BN.λ, BN.γ, BN.β λ, γ, β = BN.λ, BN.γ, BN.β
@ -125,8 +125,8 @@ function (BN::BatchNorm)(x)
m = prod(size(x)[1:end-2]) * size(x)[end] m = prod(size(x)[1:end-2]) * size(x)[end]
if !BN.active if !BN.active
μ = BN.μ μ = reshape(BN.μ, affine_shape...)
σ = BN.σ σ = reshape(BN.σ, affine_shape...)
else else
T = eltype(x) T = eltype(x)
@ -137,18 +137,18 @@ function (BN::BatchNorm)(x)
# update moving mean/std # update moving mean/std
mtm = data(convert(T, BN.momentum)) mtm = data(convert(T, BN.momentum))
BN.μ = (1 - mtm) .* BN.μ .+ mtm .* data(μ) BN.μ = (1 - mtm) .* BN.μ .+ mtm .* squeeze(data(μ), (axes...))
BN.σ = (1 - mtm) .* BN.σ .+ mtm .* data(σ) .* m ./ (m - 1) BN.σ = (1 - mtm) .* BN.σ .+ mtm .* squeeze(data(σ), (axes...)) .* m ./ (m - 1)
end end
λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ σ) .+ reshape(β, affine_shape...)) λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ σ) .+ reshape(β, affine_shape...))
end end
children(BN::BatchNorm) = children(BN::BatchNorm) =
(BN.λ, BN.β, BN.γ, BN.μ, BN.σ, BN.momentum, BN.ϵ, 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.γ), BN.μ, BN.σ, BN.momentum, BN.ϵ, 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)