def wrapped(*args, **kwargs):
start_time = time()
ret = func(*args, **kwargs)
end_time = time()
logger.info('collapsed time of calling %s: %0.4f seconds',
func.__name__, end_time - start_time)
return ret
def _tf_save_data(self, path, value):
np_array = value.np_array
if np_array.shape == ():
np_array = np.array([np_array])
with open(path, "wb") as fid:
idx2np.convert_to_file(fid, np_array)
logger.info("saving %s", path)
def transform(ugraph):
if self.APPLICABLE_LIBS is not GENERIC_SENTINEL and ugraph.lib_name not in self.APPLICABLE_LIBS:
logger.info(
"%s is not applicable to ugraph with lib name %s, skipping",
self,
ugraph.lib_name,
)
return ugraph
new_ugraph = ori_transform(ugraph)
topologic_order_graph(new_ugraph)
if self.prune_graph:
return _prune_graph(new_ugraph)
return new_ugraph
def save_idx(arr, fname):
if arr.shape == ():
arr = np.array([arr], dtype=arr.dtype)
if arr.dtype in [np.int64]:
logger.warning(
"unsupported int format for idx detected: %s, using int32 instead",
arr.dtype)
arr = arr.astype(np.int32)
out_dir = os.path.dirname(fname)
if out_dir and not os.path.exists(out_dir):
os.makedirs(out_dir)
with open(fname, "wb") as fid:
idx2np.convert_to_file(fid, arr)
logger.info("%s saved", fname)
def transform(self, ugraph):
logger.info("Transforming graph: %s", ugraph.name)
logger.info("Transform pipeline: %s", ' -> '.join(self.trans_methods))
if not self._check_generic(ugraph):
raise ValueError(
'the given graph is not generic:\n{}'.format(ugraph))
new_ugraph = self.transformer.transform(ugraph)
new_ugraph.name = ugraph.name
logger.info('Graph transormation done')
if self.save_graph:
logger.info('Saving transformed graph')
pkl_fname = "{}_transformed.pkl".format(ugraph.name)
with open(pkl_fname, 'wb') as fid:
pickle.dump(new_ugraph, fid)
logger.info('{} saved'.format(pkl_fname))
return new_ugraph
def convert_graph(model_file, output_nodes=None, config='utensor_cli.toml', target='utensor', model_name=None):
from utensor_cgen.frontend import FrontendSelector
if os.path.exists(config):
logger.info('config file {} found, reading configurations'.format(config))
with open(config) as fid:
config = loads(fid.read())
else:
config = {}
ugraph = FrontendSelector.parse(
model_file, output_nodes,
config=config,
model_name=model_name
)
backend = BackendManager.get_backend(target)(config)
backend.apply(ugraph)
return ugraph
def viz_memalloc(
cls,
ugraph,
split_on_large_graph=True,
num_tensors_per_split=20,
figsize=None,
fontsize=12,
lw=12,
cmap=_cm.BrBG_r,
rand_seed=1111
):
seed(rand_seed)
if TensorAllocationPlanner.KWARGS_NAMESCOPE not in ugraph.attributes:
logger.info('No tensor allocation plan to visualize: %s', ugraph.name)
return plt.Figure()
alloc_plan = ugraph.attributes[TensorAllocationPlanner.KWARGS_NAMESCOPE]
topo_tensors = sorted(
[tensor_name for tensor_name in alloc_plan.plan],
key=lambda tensor_name: alloc_plan.plan[tensor_name].time_slot_start
)
return cls._draw_figs(topo_tensors, alloc_plan, cmap, figsize, fontsize, lw, split_on_large_graph, num_tensors_per_split)
def apply(self, ugraph):
ref_cnts = Counter()
life_span = defaultdict(lambda: [None, None])
for op_info in ugraph.ops_info.values():
for tensor in op_info.input_tensors:
ref_cnts[tensor] += 1
for time_slot, op_name in enumerate(ugraph.topo_order):
op_info = ugraph.ops_info[op_name]
for tensor in op_info.output_tensors:
life_span[tensor][0] = time_slot
for tensor in op_info.input_tensors:
ref_cnts[tensor] -= 1
if ref_cnts[tensor] == 0:
life_span[tensor][1] = time_slot
time_alloc_plan = {}
for tensor_info, (start, end) in life_span.items():
time_alloc = TimeslotAllocation(time_slot_start=start,
time_slot_end=end)
time_alloc_plan[tensor_info] = time_alloc
logger.info('topo ordered tensor life span analysis done')
ugraph.attributes[self.KWARGS_NAMESCOPE] = time_alloc_plan
def _solve_space_alloc(self, tensors_to_schedule, nonoverlap_map):
model = cp_model.CpModel()
inter_vars = {}
tensor_allocs = {}
for tensor in tensors_to_schedule:
var_start = model.NewIntVar(0, self.max_pool_size,
'{}_start'.format(tensor.name))
var_end = model.NewIntVar(0, self.max_pool_size,
'{}_end'.format(tensor.name))
size = self._compute_tensor_bytes_size(tensor)
intv_var = model.NewIntervalVar(var_start, size, var_end,
'{}_alloc'.format(tensor.name))
inter_vars[tensor] = intv_var
tensor_allocs[tensor] = _VarMemorySpan(var_start, var_end, size)
for tensor in tensors_to_schedule:
inter_var = inter_vars[tensor]
nonoverlap_vars = [inter_vars[t] for t in nonoverlap_map[tensor]]
for other in nonoverlap_vars:
model.AddNoOverlap([inter_var, other])
var_mempool_size = model.NewIntVar(0, self.max_pool_size,
'mempool_size')
model.AddMaxEquality(var_mempool_size,
[alloc.end for alloc in tensor_allocs.values()])
model.Minimize(var_mempool_size)
solver = cp_model.CpSolver()
status = solver.Solve(model)
alloc_plan = {}
opt_mempool_size = None
if status == cp_model.OPTIMAL:
opt_mempool_size = solver.Value(var_mempool_size)
for entity, alloc in tensor_allocs.items():
alloc_plan[entity] = SpaceAllocation(
offset_start=solver.Value(alloc.start),
size=alloc.size,
data_alignment=self.data_alignment,
)
logger.info(
'optimal tensor allocation plan solved, total memory required: %i bytes',
opt_mempool_size)
logger.info('number of tensors allocated: %i' % len(alloc_plan))
else:
logger.info('tensor allocation plan not found, status: %s',
solver.StatusName(status))
if status == cp_model.INFEASIBLE:
logger.info(
'the optimal plan is infeasible, please set `max_pool_size` a larger value: %s'
% self.max_pool_size)
return alloc_plan, opt_mempool_size
def _naive_generate_files(
self, ugraph, required_ops, placeholders, tensor_var_map,
):
(
ops_map, # dict, op_info -> variable name of op in the output files
out_tensor_var_names, # list of tensor variable names, which are the variable names of the output tensors of the graph
declare_global_snippets, # list of Snippet objects, which will rendered in global scop
declare_local_snippets, # list of Snippet objects, which will rendered in local function scope
weight_snippets, # snippets for generating weights header file
) = self._get_declare_snippets(ugraph, required_ops, tensor_var_map)
# eval_snippets: List of snippet objects, which will render code snippets for tensor evaluation
eval_snippets = self._get_evaluation_snippets(ugraph, ops_map, tensor_var_map)
template_vars = {}
template_vars['model_name'] = ugraph.name
template_vars['meta_data_pool_size'] = self._compute_meta_data_size(ugraph)
template_vars['ram_data_pool_size'] = self._compute_ram_data_size(ugraph)
template_vars['placeholders'] = placeholders
template_vars['out_tensor_var_names'] = out_tensor_var_names
params_dir = Path(self.params_dir) / ugraph.name
params_dir.mkdir(parents=True, exist_ok=True)
weight_header_fname = None
if weight_snippets:
weight_header_fname = 'params_{}.hpp'.format(ugraph.name)
with (params_dir / weight_header_fname).open('w') as fid:
fid.write("/* Auto-generated by utensor cli */\n")
weight_container = ContextGlobalArrayContainer(
snippets=weight_snippets
)
fid.write(weight_container.render())
logger.info("model parameters header file generated: %s", fid.name)
# generate the compute function
model_file_dir = Path(self.model_dir)
header_fname = self.header_fname == 'None' and '{}.hpp'.format(ugraph.name) or self.header_fname
container_snippet = SimpleContainer()
container_snippet.add_declare_global_snippets(*declare_global_snippets)
container_snippet.add_declare_local_snippets(*declare_local_snippets)
container_snippet.add_eval_snippets(*eval_snippets)
container_snippet.template_vars.update(template_vars)
(model_file_dir / ugraph.name).mkdir(parents=True, exist_ok=True)
with (model_file_dir / ugraph.name / header_fname).open('w') as fid:
fid.write("/* Auto-generated by utensor cli */\n")
template = env.get_template('snippets/rearch/simple.hpp')
fid.write(template.render(**template_vars))
container_snippet.add_header(header_fname)
logger.info("model header file generated: %s", fid.name)
if weight_header_fname:
container_snippet.add_header(weight_header_fname)
composer = Composer(snippets=[container_snippet])
src_fname = self.src_fname == 'None' and '{}.cpp'.format(ugraph.name) or self.src_fname
with (model_file_dir / ugraph.name / src_fname ).open('w') as fid:
fid.write("/* Auto-generated by utensor cli */\n")
fid.write(composer.compose())
logger.info("model cpp file generated: %s", fid.name)
def save_graph(graph, graph_name="graph", out_dir="."):
out_dir = os.path.expanduser(out_dir)
graph_fname = os.path.join(out_dir, "{}.pb".format(graph_name))
with tf.gfile.FastGFile(graph_fname, "wb") as fid:
fid.write(graph.as_graph_def().SerializeToString())
logger.info("%s saved", graph_fname)
def _time_slot_generate_files(
self, ugraph, placeholders, tensor_var_map,
weight_snippets, local_snippets, declare_op_snippets, construct_op_snippets,
input_enums, output_enums,
):
# prepare template variables
template_vars = {}
template_vars['model_name'] = ugraph.name
(template_vars['meta_data_pool_size'],
template_vars['meta_dtype']) = self._compute_meta_data_size(ugraph)
(template_vars['ram_data_pool_size'],
template_vars['ram_dtype']) = self._compute_ram_data_size(ugraph)
template_vars['placeholders'] = placeholders
template_vars['out_tensor_var_names'] = [
tensor_var_map[tensor.name] for tensor in ugraph.output_tensors
]
template_vars['input_enums'] = input_enums
template_vars['output_enums'] = output_enums
params_dir = Path(self.params_dir) / ugraph.name
params_dir.mkdir(parents=True, exist_ok=True)
params_header_fname = 'params_{}.hpp'.format(ugraph.name)
with (params_dir / params_header_fname).open('w') as fid:
fid.write("/* Auto-generated by utensor cli */\n")
weight_container = ContextGlobalArrayContainer(
snippets=weight_snippets
)
fid.write(weight_container.render())
logger.info("model parameters header file generated: %s", fid.name)
model_file_dir = Path(self.model_dir)
header_fname = self.header_fname == 'None' and '{}.hpp'.format(ugraph.name) or self.header_fname
(model_file_dir / ugraph.name).mkdir(parents=True, exist_ok=True)
if self.use_model_api:
container_snippet = ModelApiContainer()
container_snippet.add_local_snippets(*local_snippets)
container_snippet.add_construct_op_snippets(*construct_op_snippets)
container_snippet.template_vars.update(template_vars)
container_snippet.add_header('"{}"'.format(params_header_fname))
with (model_file_dir / ugraph.name / header_fname).open("w") as fid:
fid.write("/* Auto-generated by utensor cli */\n")
template = env.get_template('snippets/rearch/model_api.hpp')
fid.write(
template.render(
declare_op_snippets=declare_op_snippets,
**template_vars
)
)
# container_snippet.add_header('"{}"'.format(fid.name))
container_snippet.add_header('"{}"'.format(header_fname))
logger.info("model header file generated: %s", fid.name)
else:
container_snippet = TimeSlotContainer()
container_snippet.add_local_snippets(*declare_op_snippets)
container_snippet.add_local_snippets(*local_snippets)
container_snippet.template_vars.update(template_vars)
container_snippet.add_header('"{}"'.format(params_header_fname))
with (model_file_dir / ugraph.name / header_fname).open('w') as fid:
fid.write("/* Auto-generated by utensor cli */\n")
template = env.get_template('snippets/rearch/simple.hpp')
fid.write(template.render(**template_vars))
# container_snippet.add_header('"{}"'.format(fid.name))
container_snippet.add_header('"{}"'.format(header_fname))
logger.info("model header file generated: %s", fid.name)
composer = Composer(snippets=[container_snippet])
src_fname = self.src_fname == 'None' and '{}.cpp'.format(ugraph.name) or self.src_fname
with (model_file_dir / ugraph.name / src_fname ).open('w') as fid:
fid.write("/* Auto-generated by utensor cli */\n")
fid.write(composer.compose())
logger.info("model cpp file generated: %s", fid.name)
def viz_graph(ugraph, out_fname=None, view=False, cleanup=True, colored_nodes=None, suffix=''):
if colored_nodes is None:
colored_nodes = set()
else:
colored_nodes = set(colored_nodes)
dot = Digraph()
nodes = {}
i = 0
for node in ugraph.ops:
nodes[node.name] = '{}_{}'.format(chr(ord('a') + i), suffix)
colored = node.name in colored_nodes
if colored:
dot.node(
nodes[node.name],
"%s: %s" % (node.name, node.op_type),
color='lightskyblue1',
style='filled'
)
else:
dot.node(nodes[node.name], "%s: %s" % (node.name, node.op_type))
i += 1
for n in node.input_tensors:
if n.name in nodes:
continue
nodes[n.name] = '{}_{}'.format(chr(ord('a') + i), suffix)
colored = n.name in colored_nodes
if colored:
dot.node(
nodes[n.name],
"%s: Tensor" % n.name,
color='olivedrab2',
style='filled',
shape='box',
)
else:
dot.node(nodes[n.name], "%s: Tensor" % n.name, shape='box')
i += 1
for n in node.output_tensors:
if n.name in nodes:
continue
nodes[n.name] = '{}_{}'.format(chr(ord('a') + i), suffix)
colored = n.name in colored_nodes
if colored:
dot.node(
nodes[n.name],
"%s: Tensor" % n.name,
color='olivedrab2',
style='filled',
shape='box',
)
else:
dot.node(nodes[n.name], "%s: Tensor" % n.name, shape='box')
i += 1
for node in ugraph.ops:
for n in node.input_tensors:
dot.edge(nodes[n.name], nodes[node.name])
for n in node.output_tensors:
dot.edge(nodes[node.name], nodes[n.name])
if out_fname:
dot.render(out_fname, view=view, cleanup=cleanup)
logger.info('graph visualization file generated: %s', out_fname)
return dot
def apply(self, ugraph):
if not self.enabled:
# not enabled, do nothing
return
time_alloc_plan = ugraph.attributes.get(
TopoOrderTensorTimeslotPlanner.KWARGS_NAMESCOPE)
if time_alloc_plan is None:
TopoOrderTensorTimeslotPlanner(config={}).apply(ugraph)
time_alloc_plan = ugraph.attributes[
TopoOrderTensorTimeslotPlanner.KWARGS_NAMESCOPE]
ops_to_ignore = set(ugraph.get_ops_by_type('Inline'))
if not self.include_inputs:
ops_to_ignore.update(ugraph.get_ops_by_type('Placeholder'))
if not self.include_outputs:
ops_to_ignore.update(ugraph.output_ops)
tensors_to_schedule = set()
nonoverlap_map = defaultdict(set)
for time_slot, op_name in enumerate(ugraph.topo_order):
op_info = ugraph.ops_info[op_name]
if op_info in ops_to_ignore:
continue
# all output tensor should not overlap with tensors that's still alive
for out_tensor, known_tensor in product(op_info.output_tensors,
tensors_to_schedule):
time_alloc = time_alloc_plan[known_tensor]
if time_slot in time_alloc:
nonoverlap_map[out_tensor].add(known_tensor)
# all output tensors should not overlap with each other
for out_tensor1, out_tensor2 in combinations(
op_info.output_tensors, 2):
nonoverlap_map[out_tensor1].add(out_tensor2)
# update tensors to be scheduled
tensors_to_schedule.update(op_info.output_tensors)
space_alloc_plan, opt_mempool_size = self._solve_space_alloc(
tensors_to_schedule, nonoverlap_map)
time_space_allocs = []
if space_alloc_plan:
for tensor_name in space_alloc_plan:
space_alloc = space_alloc_plan[tensor_name]
time_alloc = time_alloc_plan[tensor_name]
time_space_allocs.append(
TimeSpaceAllocation(tensor_name,
time_alloc=time_alloc,
space_alloc=space_alloc))
ugraph.attributes[self.KWARGS_NAMESCOPE] = AllocationPlan(
allocs=time_space_allocs, total_size=opt_mempool_size)
if self.out_fname:
figs = viz_memalloc(ugraph=ugraph, **self.aesthetic_kwargs)
if len(figs) == 1:
logger.info('saving tensor mem allocation to %s',
self.out_fname)
figs[0].savefig(self.out_fname)
else:
num_digits = ceil(log10(len(figs)))
file_format = '{{}}_{{:0{}d}}{{}}'.format(num_digits)
for i, fig in enumerate(figs, 1):
fname, ext = os.path.splitext(self.out_fname)
fname = file_format.format(fname, i, ext)
logger.info('saving tensor mem allocation to %s',
fname)
fig.savefig(fname)
with open('{}.pkl'.format(self.out_fname), 'wb') as fid:
pickle.dump(figs, fid)
logger.info('matplotlib figure object dumped (pickle): %s',
fid.name)
