def test_squeeze(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.squeeze(x1)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { Squeeze [op_type=Squeeze] output [op_type=Identity] input1:0 -> Squeeze '
'Squeeze:0 -> output }', onnx_to_graphviz(g))
def test_randomnormal(self):
with tf.Session() as sess:
x_ = tf.random_normal([2, 3], name="rand")
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
actual = onnx_to_graphviz(g)
expected = 'digraph { RandomNormal__2 [op_type=RandomNormal shape="[2, 3]"] output [op_type=Identity] ' \
'RandomNormal__2:0 -> output }'
self.assertEqual(expected, actual)
def test_reducesum(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.reduce_sum(x1)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { Sum [op_type=ReduceSum] output [op_type=Identity] input1:0 -> Sum Sum:0 -> output }',
onnx_to_graphviz(g))
def test_abs(self):
with tf.Session() as sess:
x = tf.placeholder(tf.float32, [2, 3], name="input")
x_ = tf.abs(x)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { Abs [op_type=Abs] output [op_type=Identity] input:0 -> Abs Abs:0 -> output }',
onnx_to_graphviz(g))
def test_reshape(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.reshape(x1, [3, 2])
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { Reshape [op_type=Reshape shape="[3, 2]"] output [op_type=Identity] input1:0 -> Reshape '
'Reshape:0 -> output }', onnx_to_graphviz(g))
def test_cast(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.cast(x1, tf.int32)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { Cast [op_type=Cast] output [op_type=Identity] input1:0 -> Cast Cast:0 -> output }',
onnx_to_graphviz(g))
def test_argminmax(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.argmin(x1, axis=0)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { ArgMin [op_type=ArgMin] output [op_type=Identity] input1:0 -> ArgMin ArgMin:0 -> output }',
onnx_to_graphviz(g))
def run_test_internal(self,
output_dict,
feed_dict,
input_names_with_port,
output_names_with_port,
rtol=0.000001):
with tf.Session() as sess:
variables_lib.global_variables_initializer().run()
expected = sess.run(output_dict, feed_dict=feed_dict)
output_name_without_port = [
n.split(':')[0] for n in output_names_with_port
]
frozen = tf.graph_util.convert_variables_to_constants(
sess, sess.graph_def, output_name_without_port)
tf.reset_default_graph()
graph_def = None
with tf.Session() as sess:
tf.import_graph_def(frozen, name='')
variables_lib.global_variables_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)
debug = False
if debug:
model_path = os.path.join(
TMPPATH, self._testMethodName + "_before_tf_optimize.pb")
with open(model_path, "wb") as f:
f.write(sess.graph_def.SerializeToString())
print("created file " + model_path)
graph_def = tf2onnx.tfonnx.tf_optimize(input_names_with_port,
output_names_with_port,
sess.graph_def, True)
if debug:
model_path = os.path.join(
TMPPATH, self._testMethodName + "_after_tf_optimize.pb")
with open(model_path, "wb") as f:
f.write(graph_def.SerializeToString())
print("created file " + model_path)
tf.reset_default_graph()
g = tf.import_graph_def(graph_def, name='')
with tf.Session(graph=g) as sess:
g = process_tf_graph(sess.graph, continue_on_error=True
) # shape_override={"output:0": [1, 6,4,5]}
actual = self._run_backend(g, output_names_with_port, feed_dict)
for expected_val, actual_val in zip(expected, actual):
self.assertAllClose(expected_val, actual_val, rtol=rtol, atol=0.)
def test_relu6(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.nn.relu6(x1)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { Relu6 [op_type=Max] Relu6__4 [op_type=Min] output [op_type=Identity] input1:0 -> Relu6 '
'Relu6__2 -> Relu6 Relu6:0 -> Relu6__4 Relu6__3 -> Relu6__4 Relu6__4:0 -> output }',
onnx_to_graphviz(g))
def test_reducesum(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.reduce_sum(x1)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph, opset=self.config.opset)
self.assertEqual(
'digraph { input1 [op_type=Placeholder shape="[2, 3]"] Sum [op_type=ReduceSum] '
'output [op_type=Identity] Const [op_type=Const] input1:0 -> Sum Sum:0 -> output }',
onnx_to_graphviz(g))
def test_rsqrt(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.rsqrt(x1)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { Rsqrt [op_type=Sqrt] Rsqrt__2 [op_type=Reciprocal] output [op_type=Identity] '
'input1:0 -> Rsqrt Rsqrt:0 -> Rsqrt__2 Rsqrt__2:0 -> output }',
onnx_to_graphviz(g))
def test_add(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x2 = tf.placeholder(tf.float32, [1, 3], name="input2")
x_ = tf.add(x1, x2)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { Add [op_type=Add] output [op_type=Identity] input1:0 -> Add input2:0 -> '
'Add Add:0 -> output }', onnx_to_graphviz(g))
def test_argminmax(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.argmin(x1, axis=0)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph, opset=self.config.opset)
self.assertEqual(
'digraph { input1 [op_type=Placeholder shape="[2, 3]"] "ArgMin/dimension" [op_type=Const] '
'ArgMin [op_type=ArgMin] output [op_type=Identity] input1:0 -> ArgMin ArgMin:0 -> output }',
onnx_to_graphviz(g))
def test_relu6(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.nn.relu6(x1)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph, opset=self.config.opset)
self.assertEqual(
'digraph { input1 [op_type=Placeholder shape="[2, 3]"] Relu6 [op_type=Relu] Relu6__2 [op_type=Clip] '
'output [op_type=Identity] input1:0 -> Relu6 Relu6:0 -> Relu6__2 Relu6__2:0 -> output }',
onnx_to_graphviz(g))
def test_squeeze(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.squeeze(x1)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph, opset=self.config.opset)
self.assertEqual(
'digraph { input1 [op_type=Placeholder shape="[2, 3]"] Squeeze [op_type=Squeeze] '
'output [op_type=Identity] input1:0 -> Squeeze Squeeze:0 -> output }',
onnx_to_graphviz(g))
def convert_tensorflow(frozen_graph_def,
name=None,
input_names=None,
output_names=None,
doc_string='',
target_opset=None,
channel_first_inputs=None,
debug_mode=False,
custom_op_conversions=None):
"""
convert a frozen tensorflow graph def into a ONNX model proto, just like how keras does.
:param frozen_graph_def: the frozen tensorflow graph
:param name: the converted onnx model internal name
:param input_names: the inputs name list of the model
:param output_names: the output name list of the model
:param doc_string: doc string
:param target_opset: the targeted onnx model opset
:param channel_first_inputs: A list of channel first input (not supported yet)
:param debug_mode: will enable the log and try to convert as much as possible on conversion
:return an ONNX ModelProto
"""
set_logger_level(logging.DEBUG if debug_mode else logging.INFO)
if target_opset is None:
target_opset = get_opset_number_from_onnx()
if not doc_string:
doc_string = "converted from {}".format(name)
graph_def = tfonnx.tf_optimize(input_names, output_names, frozen_graph_def,
True)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
if get_tensorboard_writer() is not None:
get_tensorboard_writer().add_graph(tf_graph)
custom_op_handlers = tf2onnx_builtin_conversion(target_opset)
if custom_op_conversions:
custom_op_handlers.update(custom_op_conversions)
with tf.Session(graph=tf_graph):
if not input_names:
input_nodes = list(_collect_input_nodes(tf_graph, output_names)[0])
input_names = [nd_.outputs[0].name for nd_ in input_nodes]
g = tfonnx.process_tf_graph(tf_graph,
continue_on_error=debug_mode,
opset=target_opset,
custom_op_handlers=custom_op_handlers,
inputs_as_nchw=channel_first_inputs,
output_names=output_names,
input_names=input_names)
onnx_graph = tf2onnx.optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model(doc_string)
return model_proto
def to_onnx(tf_graph,
opset=None,
shape_override=None,
middle_input_names=None):
"""Convert graph to tensorflow."""
print("-------------to_onnx--------------:", middle_input_names)
return process_tf_graph(tf_graph,
continue_on_error=False,
opset=opset,
shape_override=shape_override,
middle_inputs=middle_input_names)
def main():
args = get_args()
opset = tf2onnx.utils.find_opset(args.opset)
print("using tensorflow={}, onnx={}, opset={}, tfonnx={}/{}".format(
tf.__version__, onnx.__version__, opset, tf2onnx.__version__,
tf2onnx.version.git_version[:6]))
# override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
# support unknown dimensions.
tf2onnx.utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
if args.custom_ops:
# default custom ops for tensorflow-onnx are in the "tf" namespace
custom_ops = {
op: default_custom_op_handler
for op in args.custom_ops.split(",")
}
extra_opset = [helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)]
else:
custom_ops = {}
extra_opset = None
graph_def = tf.GraphDef()
with tf.gfile.GFile(args.input, 'rb') as f:
graph_def.ParseFromString(f.read())
# todo: consider to enable const folding by default?
graph_def = tf_optimize(args.inputs, args.outputs, graph_def,
args.fold_const)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=args.continue_on_error,
verbose=args.verbose,
target=args.target,
opset=args.opset,
custom_op_handlers=custom_ops,
extra_opset=extra_opset,
shape_override=args.shape_override,
output_names=args.outputs,
inputs_as_nchw=args.inputs_as_nchw)
optimizer = TransposeOptimizer(g, args.outputs, args.verbose is not None)
optimizer.optimize()
model_proto = g.make_model("converted from {}".format(args.input),
optimize=not args.continue_on_error)
# write onnx graph
if args.output:
with open(args.output, "wb") as f:
f.write(model_proto.SerializeToString())
def test_reshape(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x_ = tf.reshape(x1, [3, 2])
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph, opset=self.config.opset)
self.assertEqual(
'digraph { input1 [op_type=Placeholder shape="[2, 3]"] "Reshape/shape" [op_type=Const] '
'Reshape__2 [op_type=Cast] Reshape [op_type=Reshape] output [op_type=Identity] '
'"Reshape/shape":0 -> Reshape__2 input1:0 -> Reshape Reshape__2:0 -> Reshape Reshape:0 -> output }',
onnx_to_graphviz(g))
def sess2onnx(sess,
out_node_name=OUTPUT_NODE_NAME,
onnx_file_name=ONNX_FILE_NAME):
""" Converts the session to an onnx file and writes it to filename """
# Convert computational graph to protobuf (graph_def)
onnx_graph = tfonnx.process_tf_graph(sess.graph,
output_names=[out_node_name + ':0'])
model_proto = onnx_graph.make_model('onnx_model')
# Write model to file
with open(onnx_file_name, 'wb') as f:
f.write(model_proto.SerializeToString())
def test_squareddifference(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [1, 3], name="input1")
x2 = tf.placeholder(tf.float32, [1, 3], name="input2")
x_ = tf.squared_difference(x1, x2)
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { SquaredDifference [op_type=Sub] SquaredDifference__2 [op_type=Mul] '
'output [op_type=Identity] input1:0 -> SquaredDifference input2:0 -> SquaredDifference '
'SquaredDifference:0 -> SquaredDifference__2 SquaredDifference:0 -> SquaredDifference__2 '
'SquaredDifference__2:0 -> output }', onnx_to_graphviz(g))
def to_onnx(self, tf_graph, opset=None, extra_opset=None, shape_override=None, input_names=None,
const_node_values=None, initialized_tables=None):
"""Convert graph to tensorflow."""
if extra_opset is None:
extra_opset = []
if self.use_custom_ops:
extra_opset.append(utils.make_opsetid(constants.CONTRIB_OPS_DOMAIN, 1))
return process_tf_graph(tf_graph, continue_on_error=False, opset=opset,
extra_opset=extra_opset, target=Test.target, shape_override=shape_override,
input_names=input_names, output_names=self.output_names,
const_node_values=const_node_values,
initialized_tables=initialized_tables)
def _test_expand_dims(self, idx):
tf.reset_default_graph()
x_val = make_xval([3, 4])
x = tf.placeholder(tf.float32, shape=x_val.shape, name=_TFINPUT)
op = tf.expand_dims(x, idx)
with tf.Session() as sess:
output = tf.identity(op, name=_TFOUTPUT)
sess.run(tf.global_variables_initializer())
expected = sess.run(output, feed_dict={x: x_val})
g = process_tf_graph(sess.graph)
actual = self.validate_onnx(g, self._args1, {_INPUT: x_val},
expected)
self.assertAllClose(expected, actual)
def test_randomuniform(self):
with tf.Session() as sess:
shape = tf.constant([2, 3], name="shape")
x_ = tf.random_uniform(shape, name="rand")
x_ = tf.identity(x_, name="output1")
x_ = tf.identity(x_, name="output2")
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
self.assertEqual(
'digraph { RandomUniform__2 [op_type=RandomUniform shape="[2, 3]"] output1 [op_type=Identity] '
'output2 [op_type=Identity] output [op_type=Identity] RandomUniform__2:0 -> output1 '
'output1:0 -> output2 output2:0 -> output }',
onnx_to_graphviz(g))
def to_onnx(self,
tf_graph,
opset=None,
shape_override=None,
input_names=None):
"""Convert graph to tensorflow."""
return process_tf_graph(tf_graph,
continue_on_error=False,
verbose=True,
opset=opset,
target=Test.target,
shape_override=shape_override,
input_names=input_names,
output_names=self.output_names)
def convert_to_onnx(self, graph_def, inputs, outputs):
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf_loader.tf_session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=False,
target=",".join(constants.DEFAULT_TARGET),
opset=11,
input_names=inputs,
output_names=outputs,
inputs_as_nchw=None)
onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model("converted from {}".format(
self._tf_file))
return model_proto
def _conv_test(self, x_val, w, strides=None, padding="VALID"):
if strides is None:
strides = _STRIDE1x1
tf.reset_default_graph()
kernel = tf.constant(w, dtype=tf.float32, name='k')
with tf.Session() as sess:
x = tf.placeholder(tf.float32, shape=x_val.shape, name=_TFINPUT)
conv = tf.nn.conv2d(x, kernel, strides=strides, padding=padding)
output = tf.identity(conv, name=_TFOUTPUT)
sess.run(tf.global_variables_initializer())
expected = sess.run(output, feed_dict={x: x_val})
g = process_tf_graph(sess.graph)
actual = self.validate_onnx(g, self._args1, {_INPUT: x_val},
expected)
return expected, actual
def _make_onnx_model():
with tf.compat.v1.Session(graph=tf.Graph()) as session:
input_x = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 2, 3, 4], name='x')
input_y = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 2, 3, 4], name='y')
weight = tf.Variable(initial_value=4.2, dtype=tf.float32)
tf.multiply(input_x + input_y, weight, name='z')
session.run(weight.initializer)
frozen_graph_def = tf_loader.freeze_session(session, input_names=['x:0', 'y:0'], output_names=['z:0'])
with tf.compat.v1.Session(graph=tf.Graph()) as session:
tf.import_graph_def(frozen_graph_def, name='')
onnx_model = tfonnx.process_tf_graph(tf_graph=session.graph,
input_names=['x:0', 'y:0'],
output_names=['z:0']).make_model(graph_doc='Test onnx model')
return onnx_model
def convert_frozen_pb_to_onnx(frozen_pb_or_graph_def, opset=9, tf_graph_optimization=True, input_shape=None, name=None):
try:
from tf2onnx.tfonnx import process_tf_graph, tf_optimize
from tf2onnx import constants, loader, logging, utils, optimizer
except Exception as e:
logger.error('import tf2onnx error, "pip install tf2onnx"')
exit(0)
graph_def = frozen_pb_or_graph_def
if isinstance(frozen_pb_or_graph_def, str):
model_path = frozen_pb_or_graph_def
output_dir = model_path.replace('.pb', '.onnx')
else:
model_path = 'graphdef_buffer'
assert name, 'name should be give to export an .onnx when converting from a graphdef buffer'
output_dir = '{}.onnx'.format(name)
inputs, outputs = _auto_inputs_outputs_detect(graph_def)
shape_override = {}
if input_shape:
assert isinstance(input_shape, list), 'input_shape item need to be list, each for dims of a input tensor'
for idx, item in enumerate(input_shape):
shape_override[inputs[idx]] = item
# graph optimizatin with tf_graph_transform
if tf_graph_optimization:
graph_def = graph_optimization(graph_def)
with tf.Graph().as_default() as tf_graph:
tf.import_graph_def(graph_def, name='')
with tf.Session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=False,
target='',
opset=opset,
custom_op_handlers={},
extra_opset=[],
shape_override=shape_override,
input_names=inputs,
output_names=outputs,
inputs_as_nchw=None)
# graph optimization with onnx optimizer
onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model("converted from {}".format(model_path))
# write onnx graph
logger.info("")
logger.info("Successfully converted TensorFlow model %s to ONNX", model_path)
outputs = model_path.replace('.pb', '.onnx')
utils.save_protobuf(output_dir, model_proto)
logger.info("ONNX model is saved at %s", output_dir)
def test_dropout(self):
with tf.Session() as sess:
x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
x2 = tf.placeholder(tf.float32, [1, 3], name="input2")
prop = tf.placeholder(tf.float32, name="prob")
x_ = tf.add(x1, x2)
x_ = tf.nn.dropout(x_, prop)
x_ = tf.identity(x_, name="output1")
x_ = tf.identity(x_, name="output2")
_ = tf.identity(x_, name="output")
g = process_tf_graph(sess.graph)
actual = onnx_to_graphviz(g)
expected = 'digraph { Add [op_type=Add] Dropout__3 [op_type=Dropout] output1 [op_type=Identity] ' \
'output2 [op_type=Identity] output [op_type=Identity] input1:0 -> Add input2:0 -> ' \
'Add Add:0 -> Dropout__3 Dropout__3:0 -> output1 output1:0 -> output2 output2:0 -> output }'
self.assertEqual(expected, actual)
