def from_function(function, input_signature=None, opset=None, custom_ops=None, custom_op_handlers=None,
custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None, target=None,
large_model=False, output_path=None):
"""Returns a ONNX model_proto for a tf.function.
Args:
function: the tf.function we want to convert
input_signature: a tf.TensorSpec or a numpy array defining the shape/dtype of the input
opset: the opset to be used for the ONNX model, default is the latest
custom_ops: if a model contains ops not recognized by onnx runtime,
you can tag these ops with a custom op domain so that the
runtime can still open the model. Type is a dictionary `{op name: domain}`.
target: list of workarounds applied to help certain platforms
custom_op_handlers: dictionary of custom ops handlers
custom_rewriter: list of custom graph rewriters
extra_opset: list of extra opset's, for example the opset's used by custom ops
shape_override: dict with inputs that override the shapes given by tensorflow
inputs_as_nchw: transpose inputs in list from nchw to nhwc
large_model: use the ONNX external tensor storage format
output_path: save model to output_path
Returns:
An ONNX model_proto and an external_tensor_storage dict.
"""
if LooseVersion(tf.__version__) < "2.0":
raise NotImplementedError("from_function requires tf-2.0 or newer")
if not input_signature:
raise ValueError("from_function requires input_signature")
concrete_func = function.get_concrete_function(*input_signature)
input_names = [input_tensor.name for input_tensor in concrete_func.inputs
if input_tensor.dtype != tf.dtypes.resource]
output_names = [output_tensor.name for output_tensor in concrete_func.outputs
if output_tensor.dtype != tf.dtypes.resource]
initialized_tables = None
tensors_to_rename = tensor_names_from_structed(concrete_func, input_names, output_names)
with tf.device("/cpu:0"):
frozen_graph = tf_loader.from_function(concrete_func, input_names, output_names, large_model=large_model)
model_proto, external_tensor_storage = _convert_common(
frozen_graph,
name=concrete_func.name,
continue_on_error=True,
target=target,
opset=opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=input_names,
output_names=output_names,
inputs_as_nchw=inputs_as_nchw,
large_model=large_model,
tensors_to_rename=tensors_to_rename,
initialized_tables=initialized_tables,
output_path=output_path)
return model_proto, external_tensor_storage
def from_keras(model, input_signature=None, opset=None, custom_ops=None, custom_op_handlers=None,
custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None,
target=None, large_model=False, output_path=None):
"""Returns a ONNX model_proto for a tf.keras model.
Args:
model: the tf.keras model we want to convert
input_signature: a tf.TensorSpec or a numpy array defining the shape/dtype of the input
opset: the opset to be used for the ONNX model, default is the latest
target: list of workarounds applied to help certain platforms
custom_op_handlers: dictionary of custom ops handlers
custom_rewriter: list of custom graph rewriters
extra_opset: list of extra opset's, for example the opset's used by custom ops
shape_override: dict with inputs that override the shapes given by tensorflow
inputs_as_nchw: transpose inputs in list from nchw to nhwc
large_model: use the ONNX external tensor storage format
output_path: save model to output_path
Returns:
An ONNX model_proto and an external_tensor_storage dict.
"""
if LooseVersion(tf.__version__) < "2.0":
raise NotImplementedError("from_keras requires tf-1.15 or newer")
from tensorflow.python.keras.saving import saving_utils as _saving_utils # pylint: disable=import-outside-toplevel
# let tensorflow do the checking if model is a valid model
function = _saving_utils.trace_model_call(model, input_signature)
concrete_func = function.get_concrete_function(*input_signature)
input_names = [input_tensor.name for input_tensor in concrete_func.inputs
if input_tensor.dtype != tf.dtypes.resource]
output_names = [output_tensor.name for output_tensor in concrete_func.outputs
if output_tensor.dtype != tf.dtypes.resource]
initialized_tables = None
tensors_to_rename = tensor_names_from_structed(concrete_func, input_names, output_names)
with tf.device("/cpu:0"):
frozen_graph = tf_loader.from_function(concrete_func, input_names, output_names, large_model=large_model)
model_proto, external_tensor_storage = _convert_common(
frozen_graph,
name=model.name,
continue_on_error=True,
target=None,
opset=opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=input_names,
output_names=output_names,
inputs_as_nchw=inputs_as_nchw,
large_model=large_model,
tensors_to_rename=tensors_to_rename,
initialized_tables=initialized_tables,
output_path=output_path)
return model_proto, external_tensor_storage
def load(self, model_path: Union[str, Path], **_) -> Model:
if isinstance(model_path, Path):
model_path = model_path.as_posix()
get_model = load_from_file(model_path, "model", GET_MODEL_FN_NAME)
if get_model is None:
raise RuntimeError(f"Could not find {GET_MODEL_FN_NAME} in {model_path}")
model_args = filter_fn_args(self._model_args, fn=get_model)
if self._allow_growth:
physical_devices = tf.config.experimental.list_physical_devices("GPU")
for device in physical_devices:
tf.config.experimental.set_memory_growth(device, True)
tf.keras.backend.clear_session()
tf.keras.backend.set_learning_phase(False)
eager_model, call_fn = get_model(**model_args)
inputs_dict: Dict[str, TensorSpec] = {
input_name: TensorSpec(t.name, t.dtype.name, tuple(t.shape.as_list()))
for input_name, t in zip(eager_model.input_names, eager_model.inputs)
}
concrete_func = call_fn.get_concrete_function(
*[tf.TensorSpec(shape=spec.shape, dtype=spec.dtype, name=name) for name, spec in inputs_dict.items()]
)
input_tensors_names = [tensor.name for tensor in concrete_func.inputs if tensor.dtype != tf.dtypes.resource]
output_tensors_names = [tensor.name for tensor in concrete_func.outputs]
graph_def = from_function(
concrete_func, input_tensors_names, output_tensors_names, large_model=self._large_model
)
# tensor names changes after wrapping with call_fn, thus need to use those from concrete_func
outputs_dict: Dict[str, TensorSpec] = {
output_name: TensorSpec(output_tensor_name, t.dtype.name, tuple(t.shape.as_list()))
for output_name, output_tensor_name, t in zip(
eager_model.output_names, output_tensors_names, eager_model.outputs
)
}
precision = _infer_model_precision(graph_def, inputs_dict, outputs_dict)
tf.keras.backend.clear_session()
tf.keras.backend.set_learning_phase(False)
def _add_suffix_as_quickfix_for_tf24_func_refactor(spec):
if not spec.name.endswith(":0"):
spec = spec._replace(name=spec.name + ":0")
return spec
inputs_dict = {name: _add_suffix_as_quickfix_for_tf24_func_refactor(spec) for name, spec in inputs_dict.items()}
return Model(graph_def, precision, inputs_dict, outputs_dict)
def test_parse_tflite_graph(self):
def func(a, b, c):
alpha = tf.constant(1.1, dtype=tf.float32)
beta = tf.constant(2.3, dtype=tf.float32)
mul1 = tf.multiply(alpha, tf.matmul(a, b))
mul2 = tf.multiply(beta, c)
x_ = mul1 + mul2
return tf.identity(x_, name="output")
inp_shapes = [[2, 3], [3, 1], [2, 1]]
inp_dtypes = [tf.float32, tf.float32, tf.float32]
names = ['a', 'b', 'c']
names_with_port = ['a:0', 'b:0', 'c:0']
output_names = ['output']
output_names_with_port = ['output:0']
input_tensors = [tf.TensorSpec(shape=s, dtype=d, name=n) for s, d, n in zip(inp_shapes, inp_dtypes, names)]
concrete_func = tf.function(func, input_signature=tuple(input_tensors))
concrete_func = concrete_func.get_concrete_function()
graph_def = from_function(concrete_func,
input_names=names_with_port,
output_names=output_names_with_port)
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
sess_inputs = [sess.graph.get_tensor_by_name(k) for k in names_with_port]
sess_outputs = [sess.graph.get_tensor_by_name(n) for n in output_names_with_port]
converter = tf.compat.v1.lite.TFLiteConverter.from_session(sess, sess_inputs, sess_outputs)
tflite_model = converter.convert()
tflite_path = os.path.join(self.test_data_directory, self._testMethodName + ".tflite")
dir_name = os.path.dirname(tflite_path)
tflite_model = converter.convert()
os.makedirs(dir_name, exist_ok=True)
with open(tflite_path, 'wb') as f:
f.write(tflite_model)
tflite_graphs, opcodes_map, model, tensor_shapes = read_tflite_model(tflite_path)
self.assertEqual(1, len(tflite_graphs))
onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, inputs, outputs, _ = \
parse_tflite_graph(tflite_graphs[0], opcodes_map, model, tensor_shapes_override=tensor_shapes)
self.assertEqual(2, op_cnt['MUL'])
self.assertEqual(1, op_cnt['ADD'])
self.assertEqual(1, op_cnt['FULLY_CONNECTED'])
self.assertEqual(1, attr_cnt['WeightsFormat'])
self.assertEqual(names, inputs)
self.assertEqual(output_names, outputs)
for name, shape, dtype in zip(names, inp_shapes, inp_dtypes):
self.assertEqual(shape, output_shapes[name])
self.assertEqual(dtype, dtypes[name])
self.assertTrue(len(onnx_nodes) >= 4)
def run_test_case(self,
func,
feed_dict,
input_names_with_port,
output_names_with_port,
rtol=1e-07,
atol=1e-5,
convert_var_to_const=True,
constant_fold=True,
check_value=True,
check_shape=True,
check_dtype=True,
process_args=None,
onnx_feed_dict=None,
graph_validator=None,
as_session=False,
large_model=False):
# optional - passed to process_tf_graph
if process_args is None:
process_args = {}
# optional - pass distinct feed_dict to onnx runtime
if onnx_feed_dict is None:
onnx_feed_dict = feed_dict
input_names_with_port = list(feed_dict)
tf_reset_default_graph()
graph_def = None
np.random.seed(1) # Make it reproducible.
clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()}
if is_tf2() and not as_session:
#
# use eager to execute the tensorflow func
#
# numpy doesn't work for all ops, make it tf.Tensor()
input_tensors = [
tf.TensorSpec(shape=v.shape,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
input_list = [
tf.convert_to_tensor(v,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
tf.random.set_seed(1)
expected = func(*input_list)
if isinstance(expected, (list, tuple)):
# list or tuple
expected = [x.numpy() for x in expected]
else:
# single result
expected = [expected.numpy()]
# now make the eager functions a graph
concrete_func = tf.function(func,
input_signature=tuple(input_tensors))
concrete_func = concrete_func.get_concrete_function()
graph_def = from_function(concrete_func,
input_names=list(feed_dict.keys()),
output_names=output_names_with_port,
large_model=large_model)
else:
#
# use graph to execute the tensorflow func
#
with tf_session() as sess:
tf_set_random_seed(1)
input_list = []
for k, v in clean_feed_dict.items():
input_list.append(
tf_placeholder(name=k,
shape=v.shape,
dtype=tf.as_dtype(v.dtype)))
func(*input_list)
variables_lib.global_variables_initializer().run()
tf_tables_initializer().run()
output_dict = []
for out_name in output_names_with_port:
output_dict.append(sess.graph.get_tensor_by_name(out_name))
expected = sess.run(output_dict, feed_dict=feed_dict)
graph_def = freeze_session(sess,
input_names=list(feed_dict.keys()),
output_names=output_names_with_port)
tf_reset_default_graph()
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
graph_def = tf_optimize(list(feed_dict.keys()),
output_names_with_port,
graph_def,
fold_constant=constant_fold)
tf_reset_default_graph()
with tf_session() as sess:
const_node_values = None
if large_model:
const_node_values = compress_graph_def(graph_def)
tf.import_graph_def(graph_def, name='')
if self.config.is_debug_mode:
model_path = os.path.join(
self.test_data_directory,
self._testMethodName + "_after_tf_optimize.pb")
utils.save_protobuf(model_path, graph_def)
self.logger.debug("created file %s", model_path)
g = process_tf_graph(sess.graph,
opset=self.config.opset,
input_names=list(feed_dict.keys()),
output_names=output_names_with_port,
target=self.config.target,
const_node_values=const_node_values,
**process_args)
g = optimizer.optimize_graph(g)
actual = self.run_backend(g, output_names_with_port,
onnx_feed_dict, large_model)
for expected_val, actual_val in zip(expected, actual):
if check_value:
self.assertAllClose(expected_val,
actual_val,
rtol=rtol,
atol=atol)
if check_dtype:
self.assertEqual(expected_val.dtype, actual_val.dtype)
# why need shape checke: issue when compare [] with scalar
# https://github.com/numpy/numpy/issues/11071
if check_shape:
self.assertEqual(expected_val.shape, actual_val.shape)
if graph_validator:
self.assertTrue(graph_validator(g))
return g
def freeze_and_run_tf(self, func, feed_dict, outputs, as_session,
premade_placeholders, large_model, constant_fold):
np.random.seed(1) # Make it reproducible.
clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()}
if is_tf2() and not as_session:
#
# use eager to execute the tensorflow func
#
# numpy doesn't work for all ops, make it tf.Tensor()
input_tensors = [
tf.TensorSpec(shape=v.shape,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
input_list = [
tf.convert_to_tensor(v,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
tf.random.set_seed(1)
result = func(*input_list)
if isinstance(result, (list, tuple)):
# list or tuple
result = [x.numpy() for x in result]
else:
# single result
result = [result.numpy()]
# now make the eager functions a graph
concrete_func = tf.function(func,
input_signature=tuple(input_tensors))
concrete_func = concrete_func.get_concrete_function()
graph_def = from_function(concrete_func,
input_names=list(feed_dict.keys()),
output_names=outputs,
large_model=large_model)
initialized_tables = None
else:
#
# use graph to execute the tensorflow func
#
with tf_session() as sess:
tf_set_random_seed(1)
input_list = []
if not premade_placeholders:
for k, v in clean_feed_dict.items():
input_list.append(
tf_placeholder(name=k,
shape=v.shape,
dtype=tf.as_dtype(v.dtype)))
func(*input_list)
variables_lib.global_variables_initializer().run()
tf_tables_initializer().run()
output_dict = []
for out_name in outputs:
output_dict.append(sess.graph.get_tensor_by_name(out_name))
result = sess.run(output_dict, feed_dict=feed_dict)
graph_def = freeze_session(sess,
input_names=list(feed_dict.keys()),
output_names=outputs)
table_names, key_dtypes, value_dtypes = get_hash_table_info(
graph_def)
initialized_tables = {}
for n, k_dtype, val_dtype in zip(table_names, key_dtypes,
value_dtypes):
h = lookup_ops.hash_table_v2(k_dtype,
val_dtype,
shared_name=n)
k, v = lookup_ops.lookup_table_export_v2(
h, k_dtype, val_dtype)
initialized_tables[n] = (sess.run(k), sess.run(v))
tf_reset_default_graph()
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
graph_def = tf_optimize(list(feed_dict.keys()),
outputs,
graph_def,
fold_constant=constant_fold)
model_path = os.path.join(
self.test_data_directory,
self._testMethodName + "_after_tf_optimize.pb")
utils.save_protobuf(model_path, graph_def)
self.logger.debug("created file %s", model_path)
return result, graph_def, initialized_tables
def from_keras(model,
input_signature=None,
opset=None,
custom_ops=None,
custom_op_handlers=None,
custom_rewriter=None,
inputs_as_nchw=None,
extra_opset=None,
shape_override=None,
target=None,
large_model=False,
output_path=None):
"""Returns a ONNX model_proto for a tf.keras model.
Args:
model: the tf.keras model we want to convert
input_signature: a tf.TensorSpec or a numpy array defining the shape/dtype of the input
opset: the opset to be used for the ONNX model, default is the latest
target: list of workarounds applied to help certain platforms
custom_op_handlers: dictionary of custom ops handlers
custom_rewriter: list of custom graph rewriters
extra_opset: list of extra opset's, for example the opset's used by custom ops
shape_override: dict with inputs that override the shapes given by tensorflow
inputs_as_nchw: transpose inputs in list from nchw to nhwc
large_model: use the ONNX external tensor storage format
output_path: save model to output_path
Returns:
An ONNX model_proto and an external_tensor_storage dict.
"""
if LooseVersion(tf.__version__) < "2.0":
raise NotImplementedError("from_keras requires tf-2.0 or newer")
from tensorflow.python.keras.saving import saving_utils as _saving_utils # pylint: disable=import-outside-toplevel
# let tensorflow do the checking if model is a valid model
function = _saving_utils.trace_model_call(model, input_signature)
try:
concrete_func = function.get_concrete_function()
except TypeError as e:
# Legacy keras models don't accept the training arg tf provides so we hack around it
if "got an unexpected keyword argument 'training'" not in str(e):
raise e
model_call = model.call
def wrap_call(*args, training=False, **kwargs):
return model_call(*args, **kwargs)
model.call = wrap_call
function = _saving_utils.trace_model_call(model, input_signature)
concrete_func = function.get_concrete_function()
# Put it back
model.call = model_call
# These inputs will be removed during freezing (includes resources, etc.)
graph_captures = concrete_func.graph._captures # pylint: disable=protected-access
captured_inputs = [
t_name.name for t_val, t_name in graph_captures.values()
]
input_names = [
input_tensor.name for input_tensor in concrete_func.inputs
if input_tensor.name not in captured_inputs
]
output_names = [
output_tensor.name for output_tensor in concrete_func.outputs
if output_tensor.dtype != tf.dtypes.resource
]
initialized_tables = None
tensors_to_rename = tensor_names_from_structed(concrete_func, input_names,
output_names)
reverse_lookup = {v: k for k, v in tensors_to_rename.items()}
if model.output_names:
# model.output_names is an optional field of Keras models indicating output order. It is None if unused.
output_names = [reverse_lookup[out] for out in model.output_names]
elif isinstance(concrete_func.structured_outputs, dict):
# Other models specify output order using the key order of structured_outputs
output_names = [
reverse_lookup[out]
for out in concrete_func.structured_outputs.keys()
]
with tf.device("/cpu:0"):
frozen_graph = tf_loader.from_function(concrete_func,
input_names,
output_names,
large_model=large_model)
model_proto, external_tensor_storage = _convert_common(
frozen_graph,
name=model.name,
continue_on_error=True,
target=target,
opset=opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=shape_override,
input_names=input_names,
output_names=output_names,
inputs_as_nchw=inputs_as_nchw,
large_model=large_model,
tensors_to_rename=tensors_to_rename,
initialized_tables=initialized_tables,
output_path=output_path)
return model_proto, external_tensor_storage
def freeze_and_run_tf(self, func, feed_dict, outputs, as_session,
premade_placeholders, large_model):
np.random.seed(1) # Make it reproducible.
clean_feed_dict = {utils.node_name(k): v for k, v in feed_dict.items()}
if is_tf2() and not as_session:
#
# use eager to execute the tensorflow func
#
# numpy doesn't work for all ops, make it tf.Tensor()
input_tensors = [
tf.TensorSpec(shape=v.shape,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
input_list = [
tf.convert_to_tensor(v,
dtype=tf.as_dtype(v.dtype),
name=utils.node_name(k))
for k, v in feed_dict.items()
]
tf.random.set_seed(1)
result = func(*input_list)
if isinstance(result, (list, tuple)):
# list or tuple
result = [x.numpy() for x in result]
else:
# single result
result = [result.numpy()]
# now make the eager functions a graph
concrete_func = tf.function(func,
input_signature=tuple(input_tensors))
concrete_func = concrete_func.get_concrete_function()
graph_def = from_function(concrete_func,
input_names=list(feed_dict.keys()),
output_names=outputs,
large_model=large_model)
initialized_tables = None
else:
#
# use graph to execute the tensorflow func
#
with tf_session() as sess:
tf_set_random_seed(1)
input_list = []
if not premade_placeholders:
for k, v in clean_feed_dict.items():
input_list.append(
tf_placeholder(name=k,
shape=v.shape,
dtype=tf.as_dtype(v.dtype)))
func(*input_list)
variables_lib.global_variables_initializer().run()
tf_tables_initializer().run()
output_dict = []
for out_name in outputs:
output_dict.append(sess.graph.get_tensor_by_name(out_name))
result = sess.run(output_dict, feed_dict=feed_dict)
graph_def = freeze_session(sess,
input_names=list(feed_dict.keys()),
output_names=outputs)
table_info = get_hash_table_info(graph_def)
initialized_tables = {}
for info in table_info:
if info.shared_name is None:
continue
h = lookup_ops.hash_table_v2(info.key_dtype,
info.val_dtype,
shared_name=info.shared_name)
k, v = lookup_ops.lookup_table_export_v2(
h, info.key_dtype, info.val_dtype)
initialized_tables[info.shared_name] = (sess.run(k),
sess.run(v))
tf_reset_default_graph()
with tf_session() as sess:
tf.import_graph_def(graph_def, name='')
graph_def = tf_optimize(list(feed_dict.keys()), outputs,
graph_def)
return result, graph_def, initialized_tables
