|
- __precompile__(true)
-
- module PPQI
-
- using PyCall
-
- export load_config, load_model, model_forward
-
- const inference = PyNULL()
-
- function __init__()
- copy!(inference, pyimport("paddle.inference"))
- end
-
-
- struct InferenceModel
- predictor::PyObject
- input_handles::Vector{PyObject}
- output_handles::Vector{PyObject}
- end
-
-
- function load_config(modelpath::String, use_gpu::Bool=false, gpu_id::Int=0, use_mkldnn::Bool=false, cpu_treads::Int=1)::PyObject
- if isdir(modelpath)
-
- if isfile(joinpath(modelpath, "__params__"))
- model = joinpath(modelpath, "__model__")
- params = joinpath(modelpath, "__params__")
- config = inference.Config(model, params)
-
- elseif isfile(joinpath(modelpath, "params"))
- model = joinpath(modelpath, "model")
- params = joinpath(modelpath, "params")
- config = inference.Config(model, params)
-
- elseif isfile(joinpath(modelpath, "__model__"))
- config = inference.Config(modelpath)
-
- end
-
- elseif isfile(modelpath + ".pdmodel")
- model = modelpath + ".pdmodel"
- params = modelpath + ".pdiparams"
- config = inference.Config(model, params)
-
- end
-
- if use_gpu
- config.enable_use_gpu(100, gpu_id)
-
- else
- config.disable_gpu()
- config.set_cpu_math_library_num_threads(cpu_treads)
-
- if use_mkldnn
- config.enable_mkldnn()
- end
-
- end
-
- return config
-
- end
-
-
- function load_model(config::PyObject)::InferenceModel
- predictor = inference.create_predictor(config)
- input_names = predictor.get_input_names()
- output_names = predictor.get_output_names()
- input_handles = PyObject[]
-
- for input_name in input_names
- input_handle = predictor.get_input_handle(input_name)
- push!(input_handles, input_handle)
- end
-
- output_handles = PyObject[]
-
- for output_name in output_names
- output_handle = predictor.get_output_handle(output_name)
- push!(output_handles, output_handle)
- end
-
- model = InferenceModel(predictor, input_handles, output_handles)
-
- return model
- end
-
-
- function model_forward(model::InferenceModel, input_datas::Vector)::Vector
- for input_handle in model.input_handles, data in input_datas
- input_handle.copy_from_cpu(data)
- end
-
- model.predictor.run()
- output_datas = Array[]
-
- for output_handle in model.output_handles
- data = output_handle.copy_to_cpu()
- push!(output_datas, data)
- end
-
- return output_datas
-
- end
-
- end
|