def _create_handlers_variables(self, graph, vars_dict):
if self.handlers:
handlers = self.backend._get_handlers(self.opset)
for node in graph.node:
handler = handlers[node.domain].get(
node.op_type, None) if node.domain in handlers else None
if handler and bool(
handler.get_req_vars_template(node, self.initializer_dict)):
for v_name, v_template in handler.get_req_vars_template(
node, self.initializer_dict).items():
v_init, v_shape = v_template
v_name = get_variable_name(node, v_name)
if v_name in vars_dict.keys():
# found duplicated variable name due to non unique node name
exception.NON_UNIQUE_NODE_NAME_EXCEPT()
vars_dict[v_name] = tf.Variable(v_init,
dtype=v_init.dtype,
shape=v_shape,
name=v_name)
if node.op_type in ['Loop', 'Scan']:
onnx_node = OnnxNode(node)
body = onnx_node.attrs["body"]
vars_dict = self._create_handlers_variables(body, vars_dict)
elif node.op_type == 'If':
onnx_node = OnnxNode(node)
then_branch = onnx_node.attrs['then_branch']
vars_dict = self._create_handlers_variables(then_branch, vars_dict)
else_branch = onnx_node.attrs['else_branch']
vars_dict = self._create_handlers_variables(else_branch, vars_dict)
return vars_dict
def _create_handlers_variables(self, graph, vars_dict):
handlers = self.backend._get_handlers(self.opset)
for node in graph.node:
handler = handlers[node.domain].get(
node.op_type, None) if node.domain in handlers else None
if handler and bool(handler.get_req_vars_template()):
for v_name, v_template in handler.get_req_vars_template().items():
v_init, v_shape = v_template
v_count = 0
for var_name in vars_dict.keys():
v_count = v_count + 1 if var_name.startswith(v_name) else v_count
v_name = v_name + '_' + str(v_count)
vars_dict[v_name] = tf.Variable(v_init,
dtype=v_init.dtype,
shape=v_shape,
name=v_name)
if node.op_type in ['Loop', 'Scan']:
onnx_node = OnnxNode(node)
body = onnx_node.attrs["body"]
vars_dict = self._create_handlers_variables(body, vars_dict)
elif node.op_type == 'If':
onnx_node = OnnxNode(node)
then_branch = onnx_node.attrs['then_branch']
vars_dict = self._create_handlers_variables(then_branch, vars_dict)
else_branch = onnx_node.attrs['else_branch']
vars_dict = self._create_handlers_variables(else_branch, vars_dict)
return vars_dict
def get_onnx_op_from_graph_and_subgraph(graph, op_list):
for node in graph.node:
op_list[node.op_type] = 1 if node.op_type not in op_list.keys(
) else op_list[node.op_type] + 1
if node.op_type in ['Loop', 'Scan']:
onnx_node = OnnxNode(node)
body = onnx_node.attrs["body"]
op_list = get_onnx_op_from_graph_and_subgraph(body, op_list)
elif node.op_type == 'If':
onnx_node = OnnxNode(node)
then_branch = onnx_node.attrs['then_branch']
op_list = get_onnx_op_from_graph_and_subgraph(then_branch, op_list)
else_branch = onnx_node.attrs['else_branch']
op_list = get_onnx_op_from_graph_and_subgraph(else_branch, op_list)
return op_list
def onnx_graph_to_tensorflow_ops(cls, graph_def, input_values,
opset=None, strict=True):
"""
Converts ONNX graph to Tensorflow operations
Args:
graph_def: the ONNX graph to be converted
input_values: dictionary with values/tensors to initialize
the graph inputs. the dictionary must contain values
for all the graph_def.input
opset: opset version of the operator set.
strict: whether to enforce semantic equivalence between the
original model and the converted tensorflow model,
defaults to True (yes, enforce semantic equivalence).
Returns:
array of Tensorflow Tensors
"""
input_dict_items = []
# set input values for the subgraph
for value_info in graph_def.input:
if value_info.name in input_values:
x = input_values[value_info.name]
input_dict_items.append((value_info.name, x))
tensor_dict = dict(input_dict_items)
for node in graph_def.node:
onnx_node = OnnxNode(node)
output_ops = cls._onnx_node_to_tensorflow_op(onnx_node, tensor_dict,
opset=opset,strict=strict)
curr_node_output_map = \
dict(zip(onnx_node.outputs, output_ops))
tensor_dict.update(curr_node_output_map)
return tensor_dict
def onnx_graph_to_tensorflow_ops(cls,
subgraph,
tensor_dict,
opset=None,
strict=True):
"""
Converts ONNX graph to Tensorflow operations
Args:
subgraph: the ONNX graph to be converted.
tensor_dict: tensor dict of the subgraph.
opset: opset version of the operator set.
strict: whether to enforce semantic equivalence between the
original model and the converted tensorflow model,
defaults to True (yes, enforce semantic equivalence).
Returns:
array of Tensorflow Tensors
"""
for node in subgraph.node:
onnx_node = OnnxNode(node)
output_ops = cls._onnx_node_to_tensorflow_op(onnx_node,
tensor_dict,
opset=opset,
strict=strict)
curr_node_output_map = dict(zip(onnx_node.outputs, output_ops))
tensor_dict.update(curr_node_output_map)
return tensor_dict
def _get_initializer_from_graph_and_subgraphs(self, graph, graph_tensor_dict):
if graph.initializer:
graph_tensor_dict.update(
self.backend._onnx_initializer_to_input_dict_items(graph.initializer))
for node in graph.node:
if node.op_type in ['Loop', 'Scan']:
onnx_node = OnnxNode(node)
body = onnx_node.attrs["body"]
graph_tensor_dict = self._get_initializer_from_graph_and_subgraphs(
body, graph_tensor_dict)
elif node.op_type == 'If':
onnx_node = OnnxNode(node)
then_branch = onnx_node.attrs['then_branch']
graph_tensor_dict = self._get_initializer_from_graph_and_subgraphs(
then_branch, graph_tensor_dict)
else_branch = onnx_node.attrs['else_branch']
graph_tensor_dict = self._get_initializer_from_graph_and_subgraphs(
else_branch, graph_tensor_dict)
return graph_tensor_dict
def _create_handlers_variables_for_graph(cls,
handlers,
graph,
init_dict,
var_dict=None):
var_dict = dict() if var_dict is None else var_dict
for node in graph.node:
var_dict = cls._create_handler_variables_for_node(
handlers, OnnxNode(node), init_dict, var_dict)
return var_dict
def run_node(cls, node, inputs, device='CPU', outputs_info=None, **kwargs):
""" Run ONNX node.
:param node: ONNX NodeProto object.
:param inputs: Inputs.
:param device: Device run on.
:param outputs_info: None.
:param kwargs: Other args.
:return: Outputs.
"""
super(TensorflowBackend, cls).run_node(node, inputs, device)
common.sys_config.device = device
node = OnnxNode(node)
input_tensors = []
for i in inputs:
if i is None:
input_tensors.append(i)
else:
input_tensors.append(tf.constant(i))
if isinstance(inputs, dict):
feed_dict_raw = inputs
else:
assert len(node.inputs) == len(inputs)
feed_dict_raw = dict(zip(node.inputs, inputs))
# TODO: is constant the best way for feeding inputs?
input_dict = {}
for k, v in feed_dict_raw.items():
if isinstance(v, list):
list_input = []
for x in v:
if x is None:
list_input.append(x)
else:
list_input.append(tf.constant(x))
input_dict[k] = list_input
elif v is None: # keep None for empty optional data
input_dict[k] = v
else:
input_dict[k] = tf.constant(v)
module = TFModule(node, cls)
output_vals = module(**input_dict)
output_vals = [
val.numpy() if isinstance(val, tf.Tensor) else val
for val in output_vals
]
return namedtupledict('Outputs', node.outputs)(*output_vals)
def _get_initializer_from_graph_and_subgraphs(self, graph, init_dict=None):
init_dict = dict() if init_dict is None else init_dict
if graph.initializer:
init_dict.update(
self.backend._onnx_initializer_to_input_dict_items(graph.initializer))
for node in graph.node:
handler = self.handlers[node.domain].get(
node.op_type, None) if node.domain in self.handlers else None
init_dict = handler.get_initializer_from_subgraph(
OnnxNode(node), init_dict, self.
_get_initializer_from_graph_and_subgraphs) if handler else init_dict
return init_dict
def onnx_graph_to_tensorflow_ops(cls,
subgraph,
input_values,
tensor_dict,
opset=None,
strict=True):
"""
Converts ONNX graph to Tensorflow operations
Args:
subgraph: the ONNX graph to be converted
input_values: dictionary with values/tensors to initialize
the subgraph inputs. if the subgraph.input
are send in as parameters then it is required,
otherwise this can be empty dictionary
tensor_dict: the dictionary that contain values for all the
node.inputs in the subgraph that are not defined
in the subgraph or input_values.
opset: opset version of the operator set.
strict: whether to enforce semantic equivalence between the
original model and the converted tensorflow model,
defaults to True (yes, enforce semantic equivalence).
Returns:
array of Tensorflow Tensors
"""
# get the subgraph.input from input_values
subgraph_tensor_dict = input_values.copy()
# get the rest of the subgraph input from tensor_dict
for i in subgraph.input:
if i.name not in subgraph_tensor_dict.keys():
subgraph_tensor_dict[i.name] = tensor_dict[i.name]
# get the required initializer constant node(s) for the subgraph
# Need to get the initializer constant nodes from tensor_dict here
# because input from initializer will not be send in as inputs
# to the subgraph and those nodes are not in the subgraph
nodes_outputs = []
for node in subgraph.node:
for o_name in node.output:
nodes_outputs.append(o_name)
for node in subgraph.node:
for i_name in node.input:
if i_name not in nodes_outputs and i_name not in subgraph_tensor_dict.keys(
):
subgraph_tensor_dict[i_name] = tensor_dict[i_name]
onnx_node = OnnxNode(node)
output_ops = cls._onnx_node_to_tensorflow_op(onnx_node,
subgraph_tensor_dict,
opset=opset,
strict=strict)
curr_node_output_map = dict(zip(onnx_node.outputs, output_ops))
subgraph_tensor_dict.update(curr_node_output_map)
return subgraph_tensor_dict
def gen_tensor_dict(self, input_dict_items):
tensor_dict = dict(input_dict_items)
for node in self.graph_def.node:
onnx_node = OnnxNode(node)
output_ops = self.backend._onnx_node_to_tensorflow_op(onnx_node,
tensor_dict,
self.handlers,
opset=self.opset,
strict=self.strict)
curr_node_output_map = dict(zip(onnx_node.outputs, output_ops))
tensor_dict.update(curr_node_output_map)
return tensor_dict
def gen_tensor_dict(self, input_dict):
tensor_dict = self._get_initializer_from_graph_and_subgraphs(
self.graph_def, dict(input_dict))
for node in self.graph_def.node:
onnx_node = OnnxNode(node)
output_ops = self.backend._onnx_node_to_tensorflow_op(
onnx_node,
tensor_dict,
self.handlers,
opset=self.opset,
strict=self.strict)
curr_node_output_map = dict(zip(onnx_node.outputs, output_ops))
tensor_dict.update(curr_node_output_map)
return tensor_dict
def __call__(self, **kwargs):
tensor_dict = kwargs
tensor_dict.update(self.initializer_dict)
for node in self.graph_def.node:
onnx_node = OnnxNode(node)
output_ops = self.backend._onnx_node_to_tensorflow_op(onnx_node,
tensor_dict,
self.handlers,
opset=self.opset,
strict=self.strict)
curr_node_output_map = dict(zip(onnx_node.outputs, output_ops))
tensor_dict.update(curr_node_output_map)
outputs = [tensor_dict[output] for output in self.outputs]
return outputs
def run_node(cls, node, inputs, device='CPU', outputs_info=None, **kwargs):
""" Run ONNX node.
:param node: ONNX NodeProto object.
:param inputs: Inputs.
:param device: Device run on.
:param outputs_info: None.
:param kwargs: Other args.
:return: Outputs.
"""
class TFModule(tf.Module):
def __init__(self, node):
super(TFModule, self).__init__()
self.node = node
@tf.function
def __call__(self, **input_dict):
return cls._onnx_node_to_tensorflow_op(self.node, input_dict)
super(TensorflowBackend, cls).run_node(node, inputs, device)
node = OnnxNode(node)
input_tensors = []
for i in inputs:
input_tensors.append(tf.constant(i))
if isinstance(inputs, dict):
feed_dict_raw = inputs
else:
assert len(node.inputs) == len(inputs)
feed_dict_raw = dict(zip(node.inputs, inputs))
# TODO: is constant the best way for feeding inputs?
input_dict = dict([(x[0], tf.constant(x[1])) for x in feed_dict_raw.items()
])
module = TFModule(node)
output_vals = module(**input_dict)
output_vals = [
val.numpy() if isinstance(val, tf.Tensor) else val
for val in output_vals
]
return namedtupledict('Outputs', node.outputs)(*output_vals)
def gen_tensor_dict(self, input_dict):
tensor_dict = dict(input_dict)
tensor_dict.update(self.initializer_dict)
tensor_dict.update(self.handler_variables)
for node in self.graph_def.node:
onnx_node = OnnxNode(node)
output_ops = self.backend._onnx_node_to_tensorflow_op(onnx_node,
tensor_dict,
self.handlers,
opset=self.opset,
strict=self.strict)
curr_node_output_map = dict(zip(onnx_node.outputs, output_ops))
tensor_dict.update(curr_node_output_map)
# reset VAR_COUNT in handlers(currently all handlers are in ONNX_DOMAIN)
# TODO update this when we support handlers in other domain
for _, handler in self.handlers[ONNX_DOMAIN].items():
handler.VAR_COUNT = 0
return tensor_dict
def run_node(cls, node, inputs, device='CPU', outputs_info=None, **kwargs):
""" Run ONNX node.
:param node: ONNX NodeProto object.
:param inputs: Inputs.
:param device: Device run on.
:param outputs_info: None.
:param kwargs: Other args.
:return: Outputs.
"""
super(TensorflowBackend, cls).run_node(node, inputs, device)
node_graph = tf.Graph()
with node_graph.as_default():
node = OnnxNode(node)
device_option = get_device_option(Device(device))
input_tensors = []
for i in inputs:
input_tensors.append(tf.constant(i))
if isinstance(inputs, dict):
feed_dict_raw = inputs
else:
assert len(node.inputs) == len(inputs)
feed_dict_raw = dict(zip(node.inputs, inputs))
# TODO: is constant the best way for feeding inputs?
input_dict = dict([
(x[0], tf.constant(x[1])) for x in feed_dict_raw.items()
])
ops = cls._onnx_node_to_tensorflow_op(node, input_dict)
with tf.compat.v1.Session() as sess:
with tf.device(device_option):
sess.run(tf.compat.v1.global_variables_initializer())
output_vals = sess.run(ops)
return namedtupledict('Outputs', node.outputs)(*output_vals)
def _onnx_graph_to_tensorflow_rep(cls, graph_def, opset, strict):
""" Convert ONNX graph to TensorflowRep.
:param graph_def: ONNX GraphProto object.
:param opset: ONNX OperatorSetIdProto list.
:param strict: whether to enforce semantic equivalence between the original model
and the converted tensorflow model.
:return: TensorflowRep object.
"""
handlers = cls._get_handlers(opset)
tf_rep_graph = tf.Graph()
with tf_rep_graph.as_default():
# initializer: TensorProtos representing the values to initialize
# a given tensor.
# initialized: A list of names of the initialized tensors.
if graph_def.initializer:
input_dict_items = cls._onnx_initializer_to_input_dict_items(
graph_def.initializer)
initialized = {init.name for init in graph_def.initializer}
else:
input_dict_items = []
initialized = set()
# creating placeholders for currently unknown inputs
for value_info in graph_def.input:
if value_info.name in initialized:
continue
shape = list(d.dim_value if (
d.dim_value > 0 and d.dim_param == "") else None
for d in value_info.type.tensor_type.shape.dim)
value_info_name = value_info.name.replace(
":", "_tf_") + "_" + get_unique_suffix(
) if ":" in value_info.name else value_info.name
x = tf.compat.v1.placeholder(data_type.onnx2tf(
value_info.type.tensor_type.elem_type),
name=value_info_name,
shape=shape)
input_dict_items.append((value_info.name, x))
# tensor dict: this dictionary is a map from variable names
# to the latest produced TF tensors of the given name.
# This dictionary will get updated as we build the graph to
# record the names of newly produced tensors.
tensor_dict = dict(input_dict_items)
# Since tensor dict may be updated, we need to keep a copy
# of the original input dict where we track the earliest
# defined tensors so we can have access to the placeholders
# to feed in input tensors when we run the graph.
input_dict = dict(input_dict_items)
for node in graph_def.node:
onnx_node = OnnxNode(node)
output_ops = cls._onnx_node_to_tensorflow_op(onnx_node,
tensor_dict,
handlers,
opset=opset,
strict=strict)
curr_node_output_map = dict(zip(onnx_node.outputs, output_ops))
tensor_dict.update(curr_node_output_map)
tf_rep = TensorflowRep()
tf_rep.graph = tf_rep_graph
tf_rep.inputs = [
value_info.name for value_info in graph_def.input
if value_info.name not in initialized
]
tf_rep.outputs = [value_info.name for value_info in graph_def.output]
tf_rep.tensor_dict = tensor_dict
return tf_rep
def _onnx_graph_to_tensorflow_rep(cls, graph_def, opset, strict, **kwargs):
""" Convert ONNX graph to TensorflowRep.
:param graph_def: ONNX GraphProto object.
:param opset: ONNX OperatorSetIdProto list.
:param strict: whether to enforce semantic equivalence between the original model
and the converted tensorflow model.
:kwargs: additional arguements to generate tensor_dict for model debugging
:return: TensorflowRep object.
"""
# To generate tensor_dict or not, default is False
gen_tensor_dict = kwargs[
'gen_tensor_dict'] if 'gen_tensor_dict' in kwargs else False
# User provided input tensors, in the case the model inputs have unknown shapes
input_tensor_dict = kwargs[
'input_tensor_dict'] if 'input_tensor_dict' in kwargs else dict()
training_mode = kwargs[
'training_mode'] if 'training_mode' in kwargs else False
handlers = cls._get_handlers(opset)
# initializer: TensorProtos representing the values to initialize
# a given tensor.
# initialized: A list of names of the initialized tensors.
if graph_def.initializer:
initialized = {init.name for init in graph_def.initializer}
else:
initialized = set()
input_dict = dict()
module = BackendTFModule(handlers, opset, strict, graph_def, cls)
signatures = dict()
tf_rep_graph = tf.Graph()
with tf_rep_graph.as_default():
for value_info in graph_def.input:
if value_info.name in initialized:
continue
shape = list(d.dim_value if (
d.dim_value > 0 and d.dim_param == "") else None
for d in value_info.type.tensor_type.shape.dim)
value_info_name = value_info.name.replace(
":", "_tf_") + "_" + get_unique_suffix(
) if ":" in value_info.name else value_info.name
tf_spec = tf.TensorSpec(
shape,
data_type.onnx2tf(value_info.type.tensor_type.elem_type),
value_info_name)
signatures[value_info.name] = tf_spec
if gen_tensor_dict or training_mode:
x = tf.compat.v1.placeholder(
data_type.onnx2tf(
value_info.type.tensor_type.elem_type),
name=value_info_name,
shape=shape
) if value_info.name not in input_tensor_dict else input_tensor_dict[
value_info.name]
input_dict[value_info.name] = x
if gen_tensor_dict or training_mode:
input_dict_items = cls._onnx_initializer_to_input_dict_items(
graph_def.initializer, training_mode=True)
tensor_dict = dict(input_dict)
tensor_dict.update(input_dict_items)
tensor_dict[
training_flag_name] = tf.compat.v1.placeholder_with_default(
False, shape=[])
for node in graph_def.node:
onnx_node = OnnxNode(node)
output_ops = cls._onnx_node_to_tensorflow_op(onnx_node,
tensor_dict,
handlers,
opset=opset,
strict=strict)
curr_node_output_map = dict(
zip(onnx_node.outputs, output_ops))
tensor_dict.update(curr_node_output_map)
tf_rep = TensorflowRep()
tf_rep.inputs = [
value_info.name for value_info in graph_def.input
if value_info.name not in initialized
]
tf_rep.outputs = [value_info.name for value_info in graph_def.output]
module.outputs = tf_rep.outputs
tf_rep.tf_module = module
tf_rep.signatures = signatures
if gen_tensor_dict or training_mode:
tf_rep.tensor_dict = tensor_dict
if training_mode:
tf_rep.graph = tf_rep_graph
tf_rep.onnx_op_list = cls._get_onnx_op_list(graph_def)
return tf_rep
