def export_from_graph_text(
text_file,
f,
init_func=None,
constants=None,
value_info=None,
graph_name=None,
verbose=True,
):
"""Export a ONNX model from the textified ``GraphDef``.
Set the ``init_func`` to process the initializer before running.
``value_info`` should be set as (name, (onnx.TensorProto.DTYPE, shape)).
Parameters
----------
text_file : str
The path to the text file of graph def.
f : str
The file descriptor.
init_func : lambda, optional
The init function to execute before native running.
constants : dict, optional
The value of external const tensors.
value_info : dict, optional
The info external const tensors.
graph_name : str, optional
The optional graph name.
verbose : bool, optional
Whether to print the ONNX graph.
Returns
-------
onnx_pb2.ModelProto
The ONNX model.
"""
with open(text_file, 'r') as rf:
graph_def = pb.GraphDef()
parse_text_proto(rf.read(), graph_def)
export_from_graph_def(
graph_def=graph_def,
f=f,
init_func=init_func,
constants=constants,
value_info=value_info,
graph_name=graph_name,
verbose=verbose)
def import_to_graph_def(model_path):
"""Import a ONNX model to the graph def.
Parameters
----------
model_path : str
The path to the ONNX model.
Returns
-------
GraphDef
The translated graph def.
"""
if not os.path.exists(model_path):
raise ValueError('Given model({}) is not existed.'.format(model_path))
graph_def = pb.GraphDef()
serialized_proto = C.ImportONNXModel(model_path)
graph_def.ParseFromString(serialized_proto)
return graph_def
def __init__(self, name=None):
self.callback = None
self.meta_graph = _proto_def.GraphDef()
self.meta_graph.name = name if name else 'Graph'
self.graph_name = None # Determined after creating
def scan(fn, sequences, outputs_info, n_steps=None, axis=0):
"""Run a dynamic loop of the given one step function.
Parameters
----------
fn : lambda
The function to execute at each step.
sequences : Tensor or list or Tensor
The sequences.
outputs_info : Tensor or list of Tensor
The outputs.
n_steps : int or Tensor
The steps of loop.
axis : int
The axis of sequences.
Returns
-------
Tensor or list of Tensor
The outputs.
Examples
--------
>>> import dragon as dg
>>> import dragon.vm.theano as theano
>>> x = dg.Tensor('x', dtype='float32')
>>> x.set_value([1, 2, 3, 4, 5])
>>> zero = dg.Tensor('zero', dtype='float32').Variable()
>>> zero.set_value(0)
>>> prefix_sum = theano.scan(fn=lambda x_i, y_i : x_i + y_i, sequences=x, outputs_info=zero, n_steps=5, axis=0)
>>> f = theano.function(outputs=prefix_sum)
>>> print(f())
>>> [ 1. 3. 6. 10. 15.]
"""
if not isinstance(sequences, list): sequences = [sequences]
if not isinstance(outputs_info, list): outputs_info = [outputs_info]
# Extract default outputs
fn_nargs = len(inspect.getargspec(fn)[0])
default_outputs = []
for output in outputs_info:
if output is not None: default_outputs.append(output)
if len(sequences) + len(default_outputs) < fn_nargs:
raise RuntimeError('Expect {} args of fn, but at most {} args can be used\n'
'sequences provide {}, outputs_info provide {}.'. \
format(fn_nargs, len(sequences) + len(default_outputs),
len(sequences), len(default_outputs)))
# Simulate specific function
fn_inputs = [x for x in sequences] + default_outputs
fn_inputs = copy.deepcopy(fn_inputs)
# Clear to avoid importing external expressions into template function
for input in fn_inputs:
input.expressions = {}
outputs = fn(*fn_inputs)
if not isinstance(outputs, tuple): outputs = [outputs]
else: outputs = list(outputs)
if len(outputs) != len(outputs_info):
raise RuntimeError(
'Expect {} outputs of fn, but len of outputs_info is {}.'.format(
len(outputs), len(outputs_info)))
# Make GraphDef
graph_def = pb.GraphDef()
all_expressions = {}
for output in outputs:
graph_def.output.extend([output._name])
if sys.version_info >= (3, 0):
all_expressions = OrderedDict(all_expressions,
**output.expressions)
else:
all_expressions = dict(all_expressions, **output.expressions)
all_expressions = sorted(all_expressions.items(), key=lambda d: d[0])
forward_ops = copy.deepcopy([v for k, v in all_expressions])
graph_def.op.extend(forward_ops)
# Extract external inputs
external_inputs = []
internal_outputs = []
internal_inputs = [tensor.name for tensor in fn_inputs]
for op in graph_def.op:
for input in op.input:
if input not in internal_inputs:
if input not in internal_outputs:
external_inputs.append(input)
for output in op.output:
internal_outputs.append(output)
# Collect inputs (sequences + default + external)
default_outputs = [
elem.name if elem is not None else '' for elem in outputs_info
]
inputs = fn_inputs + [Tensor(name) for name in external_inputs]
kwargs = {
'axis': axis,
'nseqs': len(sequences),
'default_outputs': default_outputs,
'func_str': str(graph_def)
}
if isinstance(n_steps, int):
kwargs['nsteps'] = n_steps
kwargs['step_type'] = 'Static'
elif isinstance(n_steps, Tensor):
kwargs['extra_inputs'] = [n_steps]
kwargs['step_tensor'] = n_steps.name
kwargs['step_type'] = 'Dynamic'
else:
kwargs['step_type'] = 'Default'
kwargs['inputs_name'] = [t.name for t in inputs]
kwargs['outputs_name'] = [t.name for t in outputs]
return Tensor.CreateOperator('Scan',
inputs,
existing_outputs=outputs,
**kwargs)