Improvements for ResNet
This commit is contained in:
parent
77bb2a66de
commit
943deea92d
@ -123,9 +123,12 @@ function (BN::BatchNorm)(x)
|
|||||||
γ, β = BN.γ, BN.β
|
γ, β = BN.γ, BN.β
|
||||||
dims = length(size(x))
|
dims = length(size(x))
|
||||||
channels = size(x, dims-1)
|
channels = size(x, dims-1)
|
||||||
affine_shape = ones(Int, dims)
|
affine_shape = let dims=dims, channels=channels
|
||||||
affine_shape[end-1] = channels
|
ntuple(i->i == dims - 1 ? channels : 1, dims)
|
||||||
m = prod(size(x)[1:end-2]) * size(x)[end]
|
end
|
||||||
|
m = let sz = size(x)
|
||||||
|
prod(ntuple(i->sz[i], dims-2)) * sz[end]
|
||||||
|
end
|
||||||
|
|
||||||
if !BN.active
|
if !BN.active
|
||||||
μ = reshape(BN.μ, affine_shape...)
|
μ = reshape(BN.μ, affine_shape...)
|
||||||
|
@ -421,30 +421,3 @@ function Base.Broadcast.materialize(bc::Broadcasted{TrackedStyle})
|
|||||||
end
|
end
|
||||||
|
|
||||||
using Requires
|
using Requires
|
||||||
|
|
||||||
# https://github.com/FluxML/Flux.jl/issues/353
|
|
||||||
@init Requires.isprecompiling() || @eval Base.Broadcast begin
|
|
||||||
function flatten(bc::Broadcasted{Style}) where {Style}
|
|
||||||
isflat(bc) && return bc
|
|
||||||
args = cat_nested(bc)
|
|
||||||
let makeargs = make_makeargs(bc), f = bc.f
|
|
||||||
newf = @inline function(args::Vararg{Any,N}) where N
|
|
||||||
f(makeargs(args...)...)
|
|
||||||
end
|
|
||||||
return Broadcasted{Style}(newf, args, bc.axes)
|
|
||||||
end
|
|
||||||
end
|
|
||||||
@inline function make_makeargs(makeargs, t::Tuple{<:Broadcasted,Vararg{Any}})
|
|
||||||
bc = t[1]
|
|
||||||
let makeargs = make_makeargs(makeargs, tail(t)), f = bc.f
|
|
||||||
let makeargs = make_makeargs(makeargs, bc.args)
|
|
||||||
headargs, tailargs = make_headargs(bc.args), make_tailargs(bc.args)
|
|
||||||
return @inline function(args::Vararg{Any,N}) where N
|
|
||||||
args1 = makeargs(args...)
|
|
||||||
a, b = headargs(args1...), tailargs(args1...)
|
|
||||||
(f(a...), b...)
|
|
||||||
end
|
|
||||||
end
|
|
||||||
end
|
|
||||||
end
|
|
||||||
end
|
|
||||||
|
Loading…
Reference in New Issue
Block a user