87 lines
1.9 KiB
Julia
87 lines
1.9 KiB
Julia
import Adapt: adapt, adapt_storage
|
||
import Zygote: IdSet
|
||
|
||
children(x) = ()
|
||
mapchildren(f, x) = x
|
||
|
||
children(x::Tuple) = x
|
||
children(x::NamedTuple) = x
|
||
mapchildren(f, x::Tuple) = map(f, x)
|
||
mapchildren(f, x::NamedTuple) = map(f, x)
|
||
|
||
function treelike(m::Module, T, fs = fieldnames(T))
|
||
@eval m begin
|
||
Flux.children(x::$T) = ($([:(x.$f) for f in fs]...),)
|
||
Flux.mapchildren(f, x::$T) = $T(f.($children(x))...)
|
||
end
|
||
end
|
||
|
||
macro treelike(T, fs = nothing)
|
||
fs == nothing || isexpr(fs, :tuple) || error("@treelike T (a, b)")
|
||
fs = fs == nothing ? [] : [:($(map(QuoteNode, fs.args)...),)]
|
||
:(treelike(@__MODULE__, $(esc(T)), $(fs...)))
|
||
end
|
||
|
||
isleaf(x) = isempty(children(x))
|
||
|
||
function mapleaves(f, x; cache = IdDict())
|
||
haskey(cache, x) && return cache[x]
|
||
cache[x] = isleaf(x) ? f(x) : mapchildren(x -> mapleaves(f, x, cache = cache), x)
|
||
end
|
||
|
||
function prefor(f, x; seen = IdSet())
|
||
x ∈ seen && return
|
||
push!(seen, x)
|
||
f(x)
|
||
foreach(x -> prefor(f, x, seen = seen), children(x))
|
||
return
|
||
end
|
||
|
||
function params(m)
|
||
ps = Params()
|
||
prefor(p ->
|
||
p isa AbstractArray{<:Real} &&
|
||
!any(p′ -> p′ === p, ps) && push!(ps, p),
|
||
m)
|
||
return ps
|
||
end
|
||
|
||
params(m...) = params(m)
|
||
|
||
function loadparams!(m, xs)
|
||
for (p, x) in zip(params(m), xs)
|
||
size(p) == size(x) ||
|
||
error("Expected param size $(size(p)), got $(size(x))")
|
||
copyto!(p, x)
|
||
end
|
||
end
|
||
|
||
# CPU/GPU movement conveniences
|
||
|
||
cpu(m) = mapleaves(x -> adapt(Array, x), m)
|
||
|
||
const gpu_adaptor = if has_cuarrays()
|
||
CuArrays.cu
|
||
else
|
||
identity
|
||
end
|
||
|
||
gpu(x) = mapleaves(gpu_adaptor, x)
|
||
|
||
# Precision
|
||
|
||
adapt_storage(T::Type{<:Real}, xs::AbstractArray{<:Real}) = convert.(T, xs)
|
||
|
||
paramtype(T::Type{<:Real}, m) = mapleaves(x -> adapt(T, x), m)
|
||
|
||
f32(m) = paramtype(Float32, m)
|
||
f64(m) = paramtype(Float64, m)
|
||
|
||
# General parameter map
|
||
|
||
function mapparams(f, m)
|
||
mapleaves(m) do x
|
||
x isa Union{AbstractArray,Number} ? f(x) : x
|
||
end
|
||
end
|