Flux.jl/src/backend/tensorflow/model.jl

47 lines
1.1 KiB
Julia
Raw Normal View History

2017-05-01 16:41:42 +00:00
using Flux: mapt, collectt, shapecheckt
struct Exec
2016-10-26 13:25:10 +00:00
session::Session
input::Any
2016-10-28 20:17:48 +00:00
output::Any
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()))
2017-05-01 16:41:42 +00:00
inputs = mapt(_ -> placeholder(Float32), inputs)
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())
Exec(session, inputs, output, params, stacks)
2016-12-20 17:18:40 +00:00
end
2017-05-01 15:30:24 +00:00
retuple(xs) = xs
retuple(xs::AbstractArray{<:AbstractArray}) = (retuple.(xs)...,)
2017-05-01 16:41:42 +00:00
dictt(xs, ys) = Dict(zip(collectt(xs), collectt(ys)))
function (m::Exec)(args...)
2017-05-01 16:41:42 +00:00
shapecheckt(m.input, args)
retuple(run(m.session, m.output, dictt(m.input, args)))
2016-10-26 13:25:10 +00:00
end
mutable struct Model
model::Any
exec::Exec
Model(model) = new(model)
2016-10-26 13:25:10 +00:00
end
tf(model) = Model(model)
function (m::Model)(args...)
2017-05-01 16:41:42 +00:00
args = mapt(x->convert.(Float32, x), args)
isdefined(m, :graph) || (m.exec = makesession(m.model, 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