Flux.jl/src/compiler/loops.jl

107 lines
2.3 KiB
Julia
Raw Normal View History

2016-10-26 14:37:30 +00:00
export unroll
2016-11-07 19:44:51 +00:00
type Offset
2016-08-29 14:17:54 +00:00
name::Symbol
2016-11-07 19:44:51 +00:00
n::Int
default::Nullable{Param}
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-07 19:44:51 +00:00
function liftloops(ex, params)
2016-10-31 12:38:18 +00:00
ex = DataFlow.normedges(ex)
2016-11-07 19:44:51 +00:00
MacroTools.postwalk(ex) do ex
@capture(ex, x_{n_}) || return ex
:($(Offset(x,n))($x))
2016-08-29 14:17:54 +00:00
end
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
isa(value(v), Offset) || return v
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))
function create_steps(v::IVertex, n)
[bumpinputs(spliceinputs(v, hiddeninput(i))) for i = 1:n]
end
2016-11-08 18:02:14 +00:00
function getvar(n, step, steps, offset, default)
if step < 1
hiddeninput(sum(offset[1:n-1]) + 1 - step)
elseif step > length(steps)
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
2016-11-08 00:06:45 +00:00
function unrollgraph(v::IVertex, n)
state, offset, default = collect_state(v)
v = group(group(state...), v)
steps = create_steps(v, n)
for i = 1:n
vars = inputs(steps[i][1])
2016-11-08 18:02:14 +00:00
postwalk!(steps[i]) do v
2016-11-08 00:06:45 +00:00
isa(value(v), Offset) || return v
varid = findfirst(vars,v[1])
2016-11-08 18:02:14 +00:00
getvar(varid, value(v).n + i, steps, offset, default)
2016-11-08 00:06:45 +00:00
end
end
2016-11-08 18:02:14 +00:00
state, defaults = stateout(steps, offset, default)
group(state,group(map(x->x[2], steps)...)), map(Flux.state, defaults)
2016-11-08 00:06:45 +00:00
end
2016-11-08 18:02:14 +00:00
unrollgraph(m, n) = unrollgraph(atomise(m), n)
2016-10-25 20:10:04 +00:00
2016-10-26 10:57:03 +00:00
type Unrolled <: Model
model
graph::IVertex{Any}
2016-10-29 23:19:57 +00:00
state::Vector{Any}
2016-10-26 10:57:03 +00:00
steps::Int
end
graph(u::Unrolled) = u.graph
2016-10-28 20:17:48 +00:00
unroll(model, n) = Unrolled(model, unrollgraph(model, n)..., n)