2017-05-01 15:12:46 +00:00
|
|
|
using Flux: mapt
|
|
|
|
|
|
|
|
struct Exec
|
2016-10-26 13:25:10 +00:00
|
|
|
session::Session
|
2017-05-01 15:12:46 +00:00
|
|
|
input::Any
|
2016-10-28 20:17:48 +00:00
|
|
|
output::Any
|
2017-05-01 15:12:46 +00:00
|
|
|
params::Dict{Flux.Param,Tensor}
|
|
|
|
stacks::Dict{Any,Any}
|
2016-10-26 13:25:10 +00:00
|
|
|
end
|
|
|
|
|
2016-12-20 17:18:40 +00:00
|
|
|
function makesession(model, inputs; session = Session(Graph()))
|
2016-12-20 17:32:33 +00:00
|
|
|
params, stacks, output = tograph(model, inputs...)
|
2017-04-27 11:48:11 +00:00
|
|
|
run(session, global_variables_initializer())
|
2017-05-01 15:12:46 +00:00
|
|
|
Exec(session, inputs, output, params, stacks)
|
2016-12-20 17:18:40 +00:00
|
|
|
end
|
|
|
|
|
|
|
|
function makesession(model, n::Integer; session = Session(Graph()))
|
|
|
|
makesession(model, [placeholder(Float32) for _ = 1:n], session = session)
|
|
|
|
end
|
2016-10-28 19:50:27 +00:00
|
|
|
|
2017-05-01 15:30:24 +00:00
|
|
|
retuple(xs) = xs
|
|
|
|
retuple(xs::AbstractArray{<:AbstractArray}) = (retuple.(xs)...,)
|
|
|
|
|
2017-05-01 15:12:46 +00:00
|
|
|
function (m::Exec)(args...)
|
|
|
|
@assert length(args) == length(m.input)
|
2017-05-01 15:30:24 +00:00
|
|
|
retuple(run(m.session, m.output, Dict(zip(m.input, args))))
|
2016-10-26 13:25:10 +00:00
|
|
|
end
|
|
|
|
|
2017-05-01 15:12:46 +00:00
|
|
|
mutable struct Model
|
|
|
|
model::Any
|
|
|
|
exec::Exec
|
|
|
|
Model(model) = new(model)
|
2016-10-26 13:25:10 +00:00
|
|
|
end
|
|
|
|
|
2017-05-01 15:12:46 +00:00
|
|
|
tf(model) = Model(model)
|
|
|
|
|
|
|
|
function (m::Model)(args...)
|
|
|
|
args = mapt(x->convert.(Float32, x),args)
|
|
|
|
isdefined(m, :graph) || (m.exec = makesession(m.model, length(args)))
|
|
|
|
@tferr m.exec.stacks m.exec(args...)
|
2016-10-28 19:50:27 +00:00
|
|
|
end
|
2016-10-26 13:25:10 +00:00
|
|
|
|
2016-12-20 16:37:43 +00:00
|
|
|
for f in :[back!, update!].args
|
|
|
|
@eval function Flux.$f(m::Model, args...)
|
|
|
|
error($(string(f)) * " is not yet supported on TensorFlow models")
|
|
|
|
end
|
2016-10-26 13:25:10 +00:00
|
|
|
end
|