def GetOptimizedGraph(meta_graph):
"""Return the optimized graph.
Parameters
----------
meta_graph : dragon_pb2.GraphDef
The definition of meta graph.
Returns
-------
graph_def : dragon_pb2.GraphDef
The definition of optimized graph.
"""
from dragon.config import logger
graph_name = meta_graph.name
graph_tensor = 'GraphDef_' + graph_name
if not HasTensorCC(graph_tensor):
logger.info(
'Graph({}) does not exist, ignore printing....'.format(graph_name))
return
opt_graph_def = pb.GraphDef()
opt_graph_def.ParseFromString(FetchTensor(graph_tensor))
return opt_graph_def
def __init__(self, name=None):
self.callback = None
self.meta_graph = pb.GraphDef()
if name is None:
# Assign a auto name
self.meta_graph.name = self.graph_name = \
'Graph_' + str(ws.CURRENT_GRAPH_IDX)
ws.CURRENT_GRAPH_IDX += 1
else:
# Assign the specific name
self.meta_graph.name = self.graph_name = name
def RunGradientFlow(input_flow, targets, input_grads=None, ignored_grads=None):
"""Compute the gradients of given input flows.
Parameters
----------
input_flow : list of OperatorDef or GraphDef
The referring flows to generate gradient flows.
targets : list or str
The solving targets, generate grads automatically.
input_grads : None or list of str
The input grads.
ignored_grads : None or list of str
The grads that are explicitly ignored.
Returns
-------
None
"""
if isinstance(input_flow, list):
graph_wrapper = pb.GraphDef()
graph_wrapper.op.extend(input_flow)
input_flow = graph_wrapper
if not isinstance(input_flow, pb.GraphDef):
raise TypeError('Excepted the type of input flow is either'
'a list of OperatorDef or a GraphDef, got {}.'.format(
type(input_flow)))
from dragon.config import option, logger
log_flow = True if option['log_optimized_graph'] or option[
'log_meta_graph'] else False
RunGradientFlowCC(_stringify_proto(input_flow), targets,
input_grads if input_grads else [],
ignored_grads if ignored_grads else [],
option['share_grads'], log_flow)
if log_flow:
g_flow = pb.GraphDef()
g_flow.ParseFromString(
FetchTensor('/export/dynamic_graph/gradient_flow'))
logger.info('>>>>>>>>>>>>>>>>>> Gradient Flow <<<<<<<<<<<<<<<<<<\n')
logger.info(g_flow)
logger.info('>>>>>>>>>>>>>>>>>> Gradient Flow <<<<<<<<<<<<<<<<<<\n')
def PrintOptimizedGraph(graph_def):
graph_name = graph_def.name
graph_tensor = 'GraphDef_' + graph_name
if not HasTensorCC(graph_tensor):
logger.info(
'graph: {} does not exist, ignore printing....'.format(graph_name))
return
graph_def = pb.GraphDef()
graph_def.ParseFromString(FetchTensor(graph_tensor))
logger.info(graph_def)
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
--------
>>> x = Tensor('x', dtype='float32')
>>> x.set_value([1, 2, 3, 4, 5])
>>> zero = Tensor('zero', dtype='float32').Variable()
>>> zero.set_value(0)
>>> prefix_sum = scan(fn=lambda x_i, y_i : x_i + y_i, sequences=x,
outputs_info=zero, n_steps=x.shape[0], 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]
# 1. exact 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)))
# 2. 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)))
# 3. make GraphDef
graph_def = pb.GraphDef()
all_exprs = {}
for output in outputs:
graph_def.target.extend([output._name])
all_exprs = dict(all_exprs, **output.expressions)
all_exprs = sorted(all_exprs.items(), key=lambda d: d[0])
forward_ops = copy.deepcopy([v for k, v in all_exprs])
graph_def.op.extend(forward_ops)
# 4. exact 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)
# 5. 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(inputs,
existing_outputs=outputs,
op_type='Scan',
**kwargs)
def function(inputs=None, outputs=None, givens=None, updater=None):
"""Return a callable function that will compute ``outputs`` or apply ``updater``.
Set ``inputs`` to feed inputs into this callable function.
Set ``givens`` to substitute some tensors before making the computation graph.
Set ``updater`` to make update graph, but the update targets should be generated before.
Parameters
----------
inputs : Tensor, list of Tensor or None
The inputs to feed.
outputs : Tensor, list of Tensor or None
The outputs to solve.
givens : dict or None
The substitutions to use.
updater : BaseUpdater
The updater to use.
Returns
-------
function
The callable function.
Examples
--------
>>> x = Tensor('x').Variable()
>>> y = x * 2
>>> f = theano.function(outputs=y)
>>> x.set_value(np.ones((2, 3), dtype=np.float32))
>>> print(f())
>>> [[ 2. 2. 2.]
[ 2. 2. 2.]]
>>> f = theano.function(inputs=x, outputs=y)
>>> print(f(np.ones((2, 3), dtype=np.float32)))
>>> [[ 2. 2. 2.]
[ 2. 2. 2.]]
"""
if not isinstance(inputs, list):
if inputs is None:
inputs = []
else:
inputs = [inputs]
if not isinstance(outputs, list):
if outputs is None:
outputs = []
else:
outputs = [outputs]
if len(outputs) > 0 and updater is not None:
raise RuntimeError(
'You can specific either outputs or updater, not both.')
all_exprs = {}
all_extra_targets = set()
if not isinstance(outputs, list): outputs = [outputs]
meta_graph = pb.GraphDef()
meta_graph.name = 'Graph_' + str(ws.CURRENT_GRAPH_IDX)
ws.CURRENT_GRAPH_IDX += 1
# extract operators and targets from expressions
existing_grads = False
for output in outputs:
meta_graph.target.extend([output.name])
if sys.version_info >= (3, 0):
all_exprs = OrderedDict(all_exprs, **output.expressions)
else:
all_exprs = dict(all_exprs, **output.expressions)
all_extra_targets = all_extra_targets.union(output.extra_targets)
if len(output.grad_wrts) > 0: existing_grads = True
# we should sort out the topology of these operators before using
all_exprs = sorted(all_exprs.items(), key=lambda d: d[0])
forward_ops = copy.deepcopy([v for k, v in all_exprs])
# handle givens
if givens is not None:
name_dict = {}
external_input_exprs = {}
for old_tenosr, new_tensor in givens.items():
if isinstance(new_tensor, Tensor):
name_dict[old_tenosr.name] = new_tensor._name
if sys.version_info >= (3, 0):
external_input_exprs = OrderedDict(
external_input_exprs, **new_tensor.expressions)
else:
external_input_exprs = dict(external_input_exprs,
**new_tensor.expressions)
external_input_exprs = OrderedDict(
sorted(external_input_exprs.items(), key=lambda A: A[0]))
elif isinstance(new_tensor, np.ndarray):
ws.FeedTensor(new_tensor, GetTensorName())
all_extra_targets = all_extra_targets.union(
new_tensor.extra_targets)
external_input_ops = [v for k, v in external_input_exprs.items()]
for op in forward_ops:
op.input.extend([
name_dict[input] if input in name_dict else input
for input in op.input
])
del op.input[:int(len(op.input) / 2)]
forward_ops = external_input_ops + forward_ops
# handle grads
if existing_grads:
targets = [output.name for output in outputs]
targets.extend(all_extra_targets)
forward_ops, grad_ops = GraphGradientMaker.Make(forward_ops, targets)
else:
grad_ops = []
# Write Ops
meta_graph.op.extend(forward_ops + grad_ops)
# Write Extra Targets
for extra_target in all_extra_targets:
meta_graph.target.extend([extra_target])
# Write Misc
if len(outputs) > 0:
GraphDef_Device(meta_graph)
GraphDef_Opt(meta_graph)
GraphDef_Grad(meta_graph, outputs)
GraphDef_Phase(meta_graph, outputs)
elif updater is not None:
GraphDef_Device(meta_graph)
GraphDef_Opt(meta_graph)
GraphDef_Update(meta_graph, updater)
# call c api to create graph
ws.CreateGraph(meta_graph)
# return a lambda point to run this graph
return lambda *args, **kwargs: \
ws.RunGraph(meta_graph.name, (inputs, args), outputs, **kwargs)
def function(inputs=[], outputs=[], swaps=None, updater=None):
""" return a excutable function for a graph """
if not isinstance(inputs, list): inputs = [inputs]
if not isinstance(outputs, list): outputs = [outputs]
if len(outputs) > 0 and updater is not None:
raise RuntimeError('outputs or updater must be in 2 function.')
all_exprs = {}
all_extra_targets = set()
if not isinstance(outputs, list): outputs = [outputs]
graph_def = pb.GraphDef()
graph_def.name = 'Graph_' + str(ws.CURRENT_GRAPH_IDX)
ws.CURRENT_GRAPH_IDX += 1
# extract operators and targets from expressions
existing_grads = False
for output in outputs:
graph_def.target.extend([output.name])
if sys.version_info >= (3, 0):
all_exprs = OrderedDict(all_exprs, **output.expressions)
else:
all_exprs = dict(all_exprs, **output.expressions)
all_extra_targets = all_extra_targets.union(output.extra_targets)
if len(output.grad_wrts) > 0: existing_grads = True
for extra_target in all_extra_targets:
graph_def.target.extend([extra_target])
# we should sort out the topology of these operators before using
all_exprs = sorted(all_exprs.items(), key=lambda d: d[0])
forward_ops = copy.deepcopy([v for k, v in all_exprs])
# handle swap
if swaps is not None:
name_dict = {}
external_input_exprs = {}
for old_tenosr, new_tensor in swaps.items():
if isinstance(new_tensor, Tensor):
name_dict[old_tenosr.name] = new_tensor._name
if sys.version_info >= (3, 0):
external_input_exprs = OrderedDict(
external_input_exprs, **new_tensor.expressions)
else:
external_input_exprs = dict(external_input_exprs,
**new_tensor.expressions)
elif isinstance(new_tensor, np.ndarray):
ws.FeedTensor(new_tensor, GetTensorName())
external_input_ops = [v for k, v in external_input_exprs.items()]
for op in forward_ops:
op.input.extend([
name_dict[input] if input in name_dict else input
for input in op.input
])
del op.input[:int(len(op.input) / 2)]
forward_ops = external_input_ops + forward_ops
# handle grads
if existing_grads:
targets = [output.name for output in outputs]
forward_ops, grad_ops = GraphGradientMaker.Make(forward_ops, targets)
else:
grad_ops = []
graph_def.op.extend(forward_ops + grad_ops)
if len(outputs) > 0:
GraphDef_Device(graph_def)
GraphDef_Opt(graph_def)
GraphDef_Grad(graph_def, outputs)
GraphDef_Phase(graph_def, outputs)
elif updater is not None:
GraphDef_Device(graph_def)
GraphDef_Opt(graph_def)
GraphDef_Update(graph_def, updater)
# call c api to create graph
ws.CreateGraph(graph_def)
# return a lambda point to run this graph
return lambda *args, **kwargs: \
ws.RunGraph(graph_def.name, (inputs, args), outputs, **kwargs)
def scan(fn, sequences, outputs_info, n_steps=None, axis=0):
if not isinstance(sequences, list): sequences = [sequences]
if not isinstance(outputs_info, list): outputs_info = [outputs_info]
# 1. exact 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 {} fn args, 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)))
# 2. simulate specfic 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(
'fn expect {} outputs, but len of outputs_info is {}'.format(
len(outputs), len(outputs_info)))
# 3. make GraphDef
graph_def = pb.GraphDef()
all_exprs = {}
for output in outputs:
graph_def.target.extend([output._name])
all_exprs = dict(all_exprs, **output.expressions)
all_exprs = sorted(all_exprs.items(), key=lambda d: d[0])
forward_ops = copy.deepcopy([v for k, v in all_exprs])
graph_def.op.extend(forward_ops)
# 4. exact 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)
# 5. 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(inputs,
existing_outputs=outputs,
op_type='Scan',
**kwargs)
