stateless can be a postprocess

This commit is contained in:
Mike J Innes 2017-02-27 22:06:38 +00:00
parent 93ed44b31a
commit 69b24bfa9b

View File

@ -60,12 +60,12 @@ end
hiddeninput(n) = vertex(Split(n), inputnode(1)) hiddeninput(n) = vertex(Split(n), inputnode(1))
function create_steps(v::IVertex, n; seq = true, stateful = true) function create_steps(v::IVertex, n; seq = true)
[(stateful ? bumpinputs : copy)(seq ? spliceinputs(v, hiddeninput(i)) : v) for i = 1:n] [copy(seq ? spliceinputs(v, hiddeninput(i)) : v) for i = 1:n]
end end
function getvar(n, step, steps, offset, default; stateful = true) function getvar(n, step, steps, offset, default)
if stateful && step < 1 if step < 1
hiddeninput(sum(offset[1:n-1]) + 1 - step) hiddeninput(sum(offset[1:n-1]) + 1 - step)
elseif step 1:length(steps) elseif step 1:length(steps)
constant(default[n]) constant(default[n])
@ -84,30 +84,26 @@ function stateout(steps, offset, default)
group(outs...), defaults group(outs...), defaults
end end
function unrollgraph(v::IVertex, n; seq = true, stateful = true)
# Input: (hidden1, hidden2, ...), (x1, x2, ...) # Input: (hidden1, hidden2, ...), (x1, x2, ...)
# Output: (hidden1, hidden2, ...), (y1, y2, ...) # Output: (hidden1, hidden2, ...), (y1, y2, ...)
# If `seq` is false, takes a single `x` and uses this for each iteration. # If `seq` is false, takes a single `x` and uses this for each iteration.
# If `stateful` is false there are no hidden inputs or outputs. # If `stateful` is false there are no hidden inputs or outputs.
function unrollgraph(v::IVertex, n; seq = true)
state, offset, default = collect_state(v) state, offset, default = collect_state(v)
v = group(group(state...), v) v = group(group(state...), v)
steps = create_steps(v, n, seq = seq, stateful = stateful) steps = create_steps(v, n, seq = seq)
for i = 1:n for i = 1:n
vars = inputs(steps[i][1]) vars = inputs(steps[i][1])
postwalk!(steps[i]) do v postwalk!(steps[i]) do v
value(v) isa Offset || return v value(v) isa Offset || return v
varid = findfirst(vars,v[1]) varid = findfirst(vars,v[1])
getvar(varid, value(v).n + i, steps, offset, default, stateful = stateful) getvar(varid, value(v).n + i, steps, offset, default)
end end
end end
out = group(map(x->x[2], steps)...) out = group(map(x->x[2], steps)...)
if stateful state, defaults = stateout(steps, offset, default)
state, defaults = stateout(steps, offset, default) group(state,out), map(Flux.state, defaults)
group(state,out), map(Flux.state, defaults)
else
out, []
end
end end
unrollgraph(m, n; kws...) = unrollgraph(atomise(m), n; kws...) unrollgraph(m, n; kws...) = unrollgraph(atomise(m), n; kws...)
@ -125,9 +121,9 @@ end
graph(u::Unrolled) = u.graph graph(u::Unrolled) = u.graph
function unroll(model, n; seq = true, stateful = true) function unroll(model, n; seq = true)
graph, state = unrollgraph(model, n; seq = seq, stateful = stateful) graph, state = unrollgraph(model, n; seq = seq)
seq || stateful ? Unrolled(model, graph, state, stateful, n) : Capacitor(graph) Unrolled(model, graph, state, true, n)
end end
function unroll1(model) function unroll1(model)