def quantization_opaque_func(xgraph: XGraph, target: str, in_names: List[str],
in_tensors: List[XBuffer]) -> None:
"""
Expose quantization as an opaque function so it can be called from
both Python and C++
Arguments
---------
xgraph: XGraph
the XGraph model
target: str
the target backend for executing this xgraph
the target should be registered with a corresponding
build function.
in_names: List[str]
the names of the input names (in the same order as the input data)
in_tensors: List[XBuffer]
The input tensors (in the same order as the )
"""
def inputs_func(iter):
inputs = {
in_name: it.to_numpy()
for in_name, it in zip(in_names, in_tensors)
}
return inputs
q_xgraph = _quantize(xgraph, target, inputs_func)
xgraph.copy_from(q_xgraph)
def load_scheduled_xgraph_opaque_func(build_dir: str,
cb_scheduled_xgraph: XGraph):
"""
Expose the load scheduled xgraph function as an opaque function
so it can be called in a language agnostic way
Arguments
---------
build_dir: str
the path to the build directory containing a meta.json file
cb_scheduled_xgraph: XGraph
return the scheduled XGraph
"""
meta_file = os.path.join(build_dir, 'meta.json')
if (not os.path.isfile(meta_file)):
raise ValueError("Could not find meta file at: {}".format(meta_file))
with open(meta_file) as json_file:
meta_d = json.load(json_file)
px_net_file = meta_d['px_model']
px_params_file = meta_d['px_params']
if not os.path.isabs(px_net_file):
px_net_file = os.path.join(build_dir, px_net_file)
if not os.path.isabs(px_params_file):
px_params_file = os.path.join(build_dir, px_params_file)
scheduled_xgraph = load(px_net_file, px_params_file)
cb_scheduled_xgraph.copy_from(scheduled_xgraph)
def partition_opaque_func(xgraph: XGraph,
targets: List[str],
last_layer: str = None):
""" Expose the XGraph partition function an opaque function
so it can be called from both Python and C++ """
if last_layer == "":
last_layer = None
p_xgraph = partition(xgraph, targets, last_layer)
xgraph.copy_from(p_xgraph)
def compile_opaque_func(xgraph: XGraph, target: str,
in_tensor_names: List[str],
out_tensor_names: List[str], build_dir: str,
work_dir: str, cb_scheduled_xgraph: XGraph) -> None:
"""
Expose the compile function as an opaque function
so it can be called from both Python and C++
Arguments
---------
xgraph: XGraph
the XGraph model
target: str
the target backend for executing this xgraph
the target should be registered with a corresponding
build function.
in_tensor_names: List[str]
the names of the input tensors (in the order that they will be
provided at runtime)
out_tensor_names: List[str]
the names of the output tensors (in the order that they will be
retrieved at runtime)
build_dir: str
the directory to be used for the final build files
work_dir: str
the directory to be used for temporary work files
cb_scheduled_xgraph: XGraph
return the scheduled XGraph
"""
in_tensor_names = [stringify(itn) for itn in in_tensor_names]
out_tensor_names = [stringify(otn) for otn in out_tensor_names]
# if work_dir is None:
if not work_dir:
work_dir = os.path.abspath(os.path.join(os.getcwd(), target + "_work"))
if not build_dir:
build_dir = os.path.abspath(
os.path.join(os.getcwd(), target + "_build"))
opt_xgraph = optimize(xgraph, target)
c_xgraph = compile(opt_xgraph,
target,
work_dir=work_dir,
build_dir=build_dir)
# full_graph_input_names = xgraph.get_input_names()
# Create scheduled XGraph
# TODO: work_dir <-> build_dir
scheduled_xgraph = schedule(c_xgraph, target, work_dir=build_dir)
# Save and add to meta file
model_file = os.path.join(build_dir, 'px_model')
save(scheduled_xgraph, model_file)
meta_file = os.path.join(build_dir, 'meta.json')
if (not os.path.isfile(meta_file)):
raise ValueError("Could not find meta file at: {}".format(meta_file))
with open(meta_file, 'r') as json_file:
meta_d = json.load(json_file)
meta_d['px_model'] = 'px_model.json'
meta_d['px_params'] = 'px_model.h5'
with open(meta_file, 'w') as f:
json.dump(meta_d, f, indent=4, sort_keys=True)
# Set callback
cb_scheduled_xgraph.copy_from(scheduled_xgraph)