so can unseq

This commit is contained in:
Mike J Innes 2017-02-27 22:52:08 +00:00
parent 69b24bfa9b
commit 8f911cc31e

View File

@ -60,9 +60,7 @@ end
hiddeninput(n) = vertex(Split(n), inputnode(1)) hiddeninput(n) = vertex(Split(n), inputnode(1))
function create_steps(v::IVertex, n; seq = true) create_steps(v::IVertex, n) = [bumpinputs(spliceinputs(v, hiddeninput(i))) for i = 1:n]
[copy(seq ? spliceinputs(v, hiddeninput(i)) : v) for i = 1:n]
end
function getvar(n, step, steps, offset, default) function getvar(n, step, steps, offset, default)
if step < 1 if step < 1
@ -86,13 +84,11 @@ end
# 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 `stateful` is false there are no hidden inputs or outputs.
function unrollgraph(v::IVertex, n; seq = true) function unrollgraph(v::IVertex, n)
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) steps = create_steps(v, n)
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
@ -121,13 +117,21 @@ end
graph(u::Unrolled) = u.graph graph(u::Unrolled) = u.graph
function unroll(model, n; seq = true) function unroll(model, n)
graph, state = unrollgraph(model, n; seq = seq) graph, state = unrollgraph(model, n)
Unrolled(model, graph, state, true, n) Unrolled(model, graph, state, true, n)
end end
function unseqinput(v::IVertex)
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
function unroll1(model) function unroll1(model)
graph, state = unrollgraph(model, 1; seq = false) graph, state = unrollgraph(model, 1)
graph = unseqinput(graph)
graph = group(graph[1], map(x->x[1], inputs(graph)[2:end])...) graph = group(graph[1], map(x->x[1], inputs(graph)[2:end])...)
Unrolled(model, graph, state, false, 1) Unrolled(model, graph, state, false, 1)
end end