Flux.jl/src/compiler/loops.jl

188 lines
4.4 KiB
Julia
Raw Normal View History

2017-06-02 15:02:47 +00:00
# Stateful Models
2017-06-05 15:08:23 +00:00
mutable struct Stateful
2017-06-02 15:02:47 +00:00
model
states::Vector{Any}
istate::Vector{Any}
ostate::Vector{Any}
end
Stateful(model, ss) = Stateful(model, ss, state.(ss), state.(ss))
function (m::Stateful)(x)
m.istate = m.ostate
state, y = m.model((m.istate...,), x)
m.ostate = collect(state)
return y
end
function back!(m::Stateful, Δ, x)
back!(m.model, ((zeros.(m.ostate)...,), Δ), (m.istate...,), x)[2:end]
end
update!(m::Stateful, η) = update!(m.model, η)
2017-06-05 15:56:44 +00:00
# Seq Models
struct SeqModel
model
steps::Int
end
runseq(f, xs::Tuple...) = f(xs...)
runseq(f, xs::AbstractArray...) = stack(f(map(x -> (unstack(x,2)...,), xs)...), 2)
2017-06-05 21:49:31 +00:00
runseq(f, xs::Batch{<:Seq}...) = convert(Batch{Seq}, runseq(f, rawbatch.(xs)...))
2017-06-05 15:56:44 +00:00
runseq(f, xs) = runseq(f, (xs...,))
function (m::SeqModel)(x)
runseq(x) do x
@assert length(x) == m.steps "Expected seq length $(m.steps), got $(size(x, 2))"
m.model(x)
end
end
back!(m::SeqModel, Δ, x) = (runseq((Δ, x) -> back!(m.model, Δ, x)[1], Δ, x),)
update!(m::SeqModel, η) = update!(m.model, η)
graph(m::SeqModel) = graph(m.model)
2017-06-02 15:02:47 +00:00
# Recurrent Graphs
2017-03-14 17:56:03 +00:00
struct Offset
2016-08-29 14:17:54 +00:00
name::Symbol
2016-11-07 19:44:51 +00:00
n::Int
2017-05-30 15:53:21 +00:00
default::Nullable{Any}
2016-08-29 14:17:54 +00:00
end
2016-11-07 19:44:51 +00:00
Offset(name, n) = Offset(name, n, nothing)
2016-11-14 21:58:16 +00:00
Base.:-(o::Offset) = Offset(o.name, -o.n, o.default)
2016-11-15 20:45:24 +00:00
function liftloops(ex)
2016-10-31 12:38:18 +00:00
ex = DataFlow.normedges(ex)
2016-11-15 20:45:24 +00:00
decls = Dict()
ex = MacroTools.postwalk(ex) do ex
2016-11-07 19:44:51 +00:00
@capture(ex, x_{n_}) || return ex
2016-11-15 20:45:24 +00:00
haskey(decls, (x,n)) && return namify(decls[(x,n)])
@gensym edge
decls[(x,n)] = :($edge = $(Offset(x,n))($x))
edge
2016-08-29 14:17:54 +00:00
end
2016-11-15 20:45:24 +00:00
prepend!(ex.args, collect(values(decls)))
ex
2016-08-29 14:17:54 +00:00
end
2016-08-31 01:37:53 +00:00
function hasloops(model)
g = graph(model)
g == nothing && return false
iscyclic(g) && return true
result = false
map(m -> hasloops(m) && (result = true), g)
return result
end
function atomise(model)
postwalk(graph(model)) do v
hasloops(value(v)) || return v
spliceinputs(atomise(value(v)), inputs(v)...)
end
end
2016-08-29 14:17:54 +00:00
2016-11-08 00:06:45 +00:00
function collect_state(v::IVertex)
state = typeof(v)[]
offset = Int[]
default = Param[]
prewalk!(v) do v
2017-03-14 15:21:18 +00:00
value(v) isa Offset || return v
2016-11-08 00:06:45 +00:00
if (i = findfirst(state, v[1])) == 0
push!(state, v[1])
2016-11-08 18:02:14 +00:00
push!(offset, max(0, -value(v).n))
2016-11-08 00:06:45 +00:00
push!(default, get(value(v).default))
else
2016-11-08 18:02:14 +00:00
offset[i] = max(offset[i], -value(v).n)
2016-11-08 00:06:45 +00:00
end
v
end
return state, offset, default
end
hiddeninput(n) = vertex(Split(n), inputnode(1))
2017-02-27 22:52:08 +00:00
create_steps(v::IVertex, n) = [bumpinputs(spliceinputs(v, hiddeninput(i))) for i = 1:n]
2016-11-08 00:06:45 +00:00
2017-02-27 22:06:38 +00:00
function getvar(n, step, steps, offset, default)
if step < 1
2016-11-08 18:02:14 +00:00
hiddeninput(sum(offset[1:n-1]) + 1 - step)
2016-11-15 13:28:14 +00:00
elseif step 1:length(steps)
2016-11-08 18:02:14 +00:00
constant(default[n])
else
steps[step][1,n]
end
end
function stateout(steps, offset, default)
outs = []
defaults = []
for i = 1:length(offset), j = 1:offset[i]
push!(outs, getvar(i, length(steps)-j+1, steps, offset, default))
push!(defaults, default[i])
end
group(outs...), defaults
end
2017-02-27 21:58:08 +00:00
# Input: (hidden1, hidden2, ...), (x1, x2, ...)
# Output: (hidden1, hidden2, ...), (y1, y2, ...)
2017-06-01 17:37:20 +00:00
# TODO: make sure there's a reasonable order for hidden states
2017-02-27 21:58:08 +00:00
2017-02-27 22:52:08 +00:00
function unrollgraph(v::IVertex, n)
2016-11-08 00:06:45 +00:00
state, offset, default = collect_state(v)
v = group(group(state...), v)
2017-02-27 22:52:08 +00:00
steps = create_steps(v, n)
2016-11-08 00:06:45 +00:00
for i = 1:n
vars = inputs(steps[i][1])
2016-11-08 18:02:14 +00:00
postwalk!(steps[i]) do v
2017-03-14 15:21:18 +00:00
value(v) isa Offset || return v
2016-11-08 00:06:45 +00:00
varid = findfirst(vars,v[1])
2017-02-27 22:06:38 +00:00
getvar(varid, value(v).n + i, steps, offset, default)
2016-11-08 00:06:45 +00:00
end
end
2016-11-15 13:28:14 +00:00
out = group(map(x->x[2], steps)...)
2017-02-27 22:06:38 +00:00
state, defaults = stateout(steps, offset, default)
2017-06-02 14:46:24 +00:00
group(state,out), defaults
2016-11-08 00:06:45 +00:00
end
2016-11-15 13:28:14 +00:00
unrollgraph(m, n; kws...) = unrollgraph(atomise(m), n; kws...)
2016-10-25 20:10:04 +00:00
2017-02-27 22:52:08 +00:00
function unroll(model, n)
graph, state = unrollgraph(model, n)
2017-03-21 01:32:12 +00:00
SeqModel(Stateful(Capacitor(graph), state), n)
2016-11-15 13:28:14 +00:00
end
2016-11-14 21:58:16 +00:00
2017-06-02 15:11:09 +00:00
function stateless(s::Stateful)
v = graph(s.model)
v = spliceinputs(v, group(constant.(s.states)...),
[inputnode(i) for i = 1:graphinputs(v)-1]...)
2017-06-02 15:22:55 +00:00
Capacitor(v[2])
2017-06-02 15:11:09 +00:00
end
stateless(s::SeqModel) = SeqModel(stateless(s.model), s.steps)
2017-02-27 23:05:19 +00:00
function unseqin(v::IVertex)
2017-02-27 22:52:08 +00:00
prewalk(v) do v
# TODO: inputidx function
isa(value(v), Split) && DataFlow.isinput(v[1]) && value(v[1]).n == 2 ? v[1] : v
end
end
2017-06-02 15:22:39 +00:00
unseqout(v::IVertex) = group(v[1], v[2][1])
2017-02-27 23:05:19 +00:00
unseq(graph) = unseqout(unseqin(graph))
2016-11-15 20:01:35 +00:00
function unroll1(model)
2017-02-27 22:52:08 +00:00
graph, state = unrollgraph(model, 1)
2017-03-29 17:30:28 +00:00
Stateful(Capacitor(unseq(graph)), state)
2016-11-15 20:01:35 +00:00
end
2017-03-14 15:21:18 +00:00
flip(model) = Capacitor(map(x -> x isa Offset ? -x : x, atomise(model)))