unrolled models without state input
This commit is contained in:
parent
85335e5ed0
commit
8fceb62d33
@ -54,14 +54,14 @@ end
|
|||||||
|
|
||||||
hiddeninput(n) = vertex(Split(n), inputnode(1))
|
hiddeninput(n) = vertex(Split(n), inputnode(1))
|
||||||
|
|
||||||
function create_steps(v::IVertex, n; seq = true)
|
function create_steps(v::IVertex, n; seq = true, stateful = true)
|
||||||
[bumpinputs(seq ? spliceinputs(v, hiddeninput(i)) : v) for i = 1:n]
|
[(stateful ? bumpinputs : copy)(seq ? spliceinputs(v, hiddeninput(i)) : v) for i = 1:n]
|
||||||
end
|
end
|
||||||
|
|
||||||
function getvar(n, step, steps, offset, default)
|
function getvar(n, step, steps, offset, default; stateful = true)
|
||||||
if step < 1
|
if stateful && step < 1
|
||||||
hiddeninput(sum(offset[1:n-1]) + 1 - step)
|
hiddeninput(sum(offset[1:n-1]) + 1 - step)
|
||||||
elseif step > length(steps)
|
elseif step ∉ 1:length(steps)
|
||||||
constant(default[n])
|
constant(default[n])
|
||||||
else
|
else
|
||||||
steps[step][1,n]
|
steps[step][1,n]
|
||||||
@ -78,33 +78,43 @@ function stateout(steps, offset, default)
|
|||||||
group(outs...), defaults
|
group(outs...), defaults
|
||||||
end
|
end
|
||||||
|
|
||||||
function unrollgraph(v::IVertex, n; seq = true)
|
function unrollgraph(v::IVertex, n; seq = true, stateful = 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)
|
steps = create_steps(v, n, seq = seq, stateful = stateful)
|
||||||
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
|
||||||
isa(value(v), Offset) || return v
|
isa(value(v), Offset) || return v
|
||||||
varid = findfirst(vars,v[1])
|
varid = findfirst(vars,v[1])
|
||||||
getvar(varid, value(v).n + i, steps, offset, default)
|
getvar(varid, value(v).n + i, steps, offset, default, stateful = stateful)
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
state, defaults = stateout(steps, offset, default)
|
out = group(map(x->x[2], steps)...)
|
||||||
group(state,group(map(x->x[2], steps)...)), map(Flux.state, defaults)
|
if stateful
|
||||||
|
state, defaults = stateout(steps, offset, default)
|
||||||
|
group(state,out), map(Flux.state, defaults)
|
||||||
|
else
|
||||||
|
out, []
|
||||||
|
end
|
||||||
end
|
end
|
||||||
|
|
||||||
unrollgraph(m, n; seq = true) = unrollgraph(atomise(m), n; seq = seq)
|
unrollgraph(m, n; kws...) = unrollgraph(atomise(m), n; kws...)
|
||||||
|
|
||||||
|
# TODO: perhaps split into SeqModel + StatefulModel
|
||||||
type Unrolled <: Model
|
type Unrolled <: Model
|
||||||
model
|
model
|
||||||
graph::IVertex{Any}
|
graph::IVertex{Any}
|
||||||
state::Vector{Any}
|
state::Vector{Any}
|
||||||
|
stateful::Bool
|
||||||
steps::Int
|
steps::Int
|
||||||
end
|
end
|
||||||
|
|
||||||
graph(u::Unrolled) = u.graph
|
graph(u::Unrolled) = u.graph
|
||||||
|
|
||||||
unroll(model, n; seq = true) = Unrolled(model, unrollgraph(model, n; seq = seq)..., n)
|
function unroll(model, n; seq = true, stateful = true)
|
||||||
|
graph, state = unrollgraph(model, n; seq = seq, stateful = stateful)
|
||||||
|
seq || stateful ? Unrolled(model, graph, state, stateful, n) : Capacitor(graph)
|
||||||
|
end
|
||||||
|
|
||||||
flip(model) = Capacitor(map(x -> isa(x, Offset) ? -x : x, atomise(model)))
|
flip(model) = Capacitor(map(x -> isa(x, Offset) ? -x : x, atomise(model)))
|
||||||
|
Loading…
Reference in New Issue
Block a user