Merge branch 'master' of https://github.com/onnx/tensorflow-onnx into opset_9_scan

Author: pengwa

tests/run_pretrained_models.py

@@ -153,7 +153,7 @@ def run_tensorflow(self, sess, inputs):
 
     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=True, verbose=True, opset=opset,
+        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)
 
@@ -186,7 +186,6 @@ def run_onnxruntime(self, name, model_proto, inputs):
         """Run test against msrt-next backend."""
         import onnxruntime as rt
         model_path = utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto, include_test_data=True)
-        utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto, include_test_data=False, as_text=True)
         print("\t\t" + model_path)
         m = rt.InferenceSession(model_path)
         results = m.run(self.output_names, inputs)
@@ -266,7 +265,7 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
                 onnx_graph = self.to_onnx(sess.graph, opset=opset, shape_override=shape_override,
                                           input_names=inputs.keys())
                 model_proto = onnx_graph.make_model("converted from tf2onnx")
-                new_model_proto = GraphUtil.opt_transposes_with_graph(onnx_graph, "test", debug=debug)
+                new_model_proto = GraphUtil.optimize_graph(onnx_graph, "test", debug=debug)
                 if new_model_proto:
                     model_proto = new_model_proto
                 else:

tests/test_backend.py

@@ -620,23 +620,27 @@ def test_logicaland(self):
 
     @check_onnxruntime_incompatibility("Greater")
     def test_greater(self):
-        x_val1 = np.array([4, 2, 4, 1], dtype=np.float32).reshape((2, 2))
-        x_val2 = np.array([2, 4, 4, 1], dtype=np.float32).reshape((2, 2))
-        x1 = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT)
-        x2 = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT1)
-        mi = tf.greater(x1, x2)
-        _ = tf.identity(mi, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
+        for op in [tf.greater, tf.greater_equal]:
+            tf.reset_default_graph()
+            x_val1 = np.array([4, 2, 4, 1], dtype=np.float32).reshape((2, 2))
+            x_val2 = np.array([2, 4, 4, 1], dtype=np.float32).reshape((2, 2))
+            x1 = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT)
+            x2 = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT1)
+            mi = op(x1, x2)
+            _ = tf.identity(mi, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
 
     @check_onnxruntime_incompatibility("Greater")
     def test_greater_unsupport_type(self):
-        x_val1 = np.array([4, 2, 4, 1], dtype=np.int32).reshape((2, 2))
-        x_val2 = np.array([2, 4, 4, 1], dtype=np.int32).reshape((2, 2))
-        x1 = tf.placeholder(tf.int32, [2, 2], name=_TFINPUT)
-        x2 = tf.placeholder(tf.int32, [2, 2], name=_TFINPUT1)
-        mi = tf.greater(x1, x2)
-        _ = tf.identity(mi, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
+        for op in [tf.greater, tf.greater_equal]:
+            tf.reset_default_graph()
+            x_val1 = np.array([4, 2, 4, 1], dtype=np.int32).reshape((2, 2))
+            x_val2 = np.array([2, 4, 4, 1], dtype=np.int32).reshape((2, 2))
+            x1 = tf.placeholder(tf.int32, [2, 2], name=_TFINPUT)
+            x2 = tf.placeholder(tf.int32, [2, 2], name=_TFINPUT1)
+            mi = op(x1, x2)
+            _ = tf.identity(mi, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
 
     @check_onnxruntime_incompatibility("Less")
     def test_less(self):

tests/test_optimizers.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 from onnx import helper, TensorProto
+from tf2onnx import utils
 from tf2onnx.graph import GraphUtil
 from backend_test_base import Tf2OnnxBackendTestBase
 from common import unittest_main
@@ -19,22 +20,23 @@
 class OptimizerTests(Tf2OnnxBackendTestBase):
     """Run original model proto and modified model proto with onnxruntime, compare the results."""
 
-    def run_and_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
-                        remaining_transpose_num=None, debug=False, rtol=1e-07):
+    def run_and_compare(self, output_names_with_port, onnx_feed_dict, origin_proto, op_type,
+                        remaining_op_num, debug=False, rtol=1e-07):
+        utils.make_sure(op_type is not None, "op_type should be specified")
+        utils.make_sure(remaining_op_num is not None, "remaining_op_num should be specified")
+
         origin_model_path = self.save_onnx_model(origin_proto, onnx_feed_dict, postfix="_origin")
 
-        new_proto = GraphUtil.opt_transposes_with_model_proto(origin_proto)
+        new_proto = GraphUtil.optimize_graph_with_model_proto(origin_proto)
 
         self.assertTrue(new_proto, msg="model proto after optimizer should not be None")
 
         new_model_path = self.save_onnx_model(new_proto, onnx_feed_dict, postfix="_opt")
-
-        previous = GraphUtil.get_node_count_from_onnx_graph(origin_proto.graph)
         current = GraphUtil.get_node_count_from_onnx_graph(new_proto.graph)
 
-        self.assertTrue(current["Transpose"] < previous["Transpose"], msg="transpose ops count not changed")
-        if remaining_transpose_num is not None:
-            self.assertTrue(current["Transpose"] == remaining_transpose_num, msg="some transpose ops left unexpected")
+        self.assertTrue(current[op_type] == remaining_op_num,
+                        msg="Expect " + str(remaining_op_num) + " " + op_type + " ops left, but actually " +
+                        str(current[op_type]) + " left")
 
         if self.config.is_onnxruntime_backend:
             expected = self.run_onnxruntime(origin_model_path, onnx_feed_dict, output_names_with_port)
@@ -47,7 +49,14 @@ def run_and_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
             self.assertEqual(expected_val.dtype, actual_val.dtype)
             self.assertEqual(expected_val.shape, actual_val.shape)
 
-    def test_relu(self):
+    # Tranpose Optimizer Tests Start
+
+    def run_transpose_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
+                              remaining_transpose_num=None, debug=False, rtol=1e-07):
+        self.run_and_compare(output_names_with_port, onnx_feed_dict, origin_proto, op_type="Transpose",
+                             remaining_op_num=remaining_transpose_num, debug=debug, rtol=rtol)
+
+    def test_transpose_relu(self):
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans_1")
         node2 = helper.make_node("Relu", ["Y"], ["Z"], name="relu")
         node3 = helper.make_node("Transpose", ["Z"], ["Z1"], perm=[0, 3, 1, 2], name="trans_2")
@@ -60,10 +69,10 @@ def test_relu(self):
         )
 
         model_proto = helper.make_model(graph, producer_name="onnx-tests")
-        self.run_and_compare(["Z1"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
-                             model_proto, remaining_transpose_num=0)
+        self.run_transpose_compare(["Z1"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
 
-    def test_leaky_relu(self):
+    def test_transpose_leaky_relu(self):
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans_1")
         node2 = helper.make_node("LeakyRelu", ["Y"], ["Z"], alpha=0.02, name="relu")
         node3 = helper.make_node("Transpose", ["Z"], ["Z1"], perm=[0, 3, 1, 2], name="trans_2")
@@ -76,10 +85,10 @@ def test_leaky_relu(self):
         )
 
         model_proto = helper.make_model(graph, producer_name="onnx-tests")
-        self.run_and_compare(["Z1"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
-                             model_proto, remaining_transpose_num=0)
+        self.run_transpose_compare(["Z1"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
 
-    def test_max(self):
+    def test_transpose_max(self):
         const_1_val = [2.0]
         const_1 = helper.make_tensor("const_1", TensorProto.FLOAT, (1,), const_1_val)
         const_1_node = helper.make_node("Constant", [], ["const_1"], value=const_1, name="const_1")
@@ -104,8 +113,8 @@ def test_max(self):
         )
 
         model_proto = helper.make_model(graph, producer_name="onnx-tests")
-        self.run_and_compare(["Z1"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
-                             model_proto, remaining_transpose_num=0)
+        self.run_transpose_compare(["Z1"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
 
     def test_transpose_merge(self):
         node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
@@ -120,8 +129,8 @@ def test_transpose_merge(self):
         )
 
         model_proto = helper.make_model(graph, producer_name="onnx-tests")
-        self.run_and_compare(["OUT"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
-                             model_proto, remaining_transpose_num=1)
+        self.run_transpose_compare(["OUT"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=1)
 
     def test_transpose_with_shape(self):
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
@@ -135,9 +144,135 @@ def test_transpose_with_shape(self):
         )
 
         model_proto = helper.make_model(graph, producer_name="onnx-tests")
-        self.run_and_compare(["Z"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
-                             model_proto, remaining_transpose_num=0)
+        self.run_transpose_compare(["Z"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
+    # Tranpose Optimizer Tests End
+
+    # Identity Optimizer Tests Start
+
+    def run_identity_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
+                             remaining_identity_num=None, debug=False, rtol=1e-07):
+        self.run_and_compare(output_names_with_port, onnx_feed_dict, origin_proto, op_type="Identity",
+                             remaining_op_num=remaining_identity_num, debug=debug, rtol=rtol)
+
+    def test_identity_non_graph_output(self):
+        node1 = helper.make_node("Add", ["X", "X"], ["Y"], name="add")
+        node2 = helper.make_node("Identity", ["Y"], ["Z"], name="identity")
+        node3 = helper.make_node("Shape", ["Z"], ["Z1"], name="shape")
+
+        graph = helper.make_graph(
+            [node1, node2, node3],
+            "identity-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4, 5))],
+            [helper.make_tensor_value_info("Z1", TensorProto.INT64, [4])],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_identity_compare(["Z1"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                  model_proto, remaining_identity_num=0)
+
+    def test_identity_unremovable_identity(self):
+        # should not remove!!
+        node1 = helper.make_node("Identity", ["X"], ["Y"], name="identity")
+
+        graph = helper.make_graph(
+            [node1],
+            "identity-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4, 5))],
+            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (2, 3, 4, 5))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_identity_compare(["Y"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                  model_proto, remaining_identity_num=1)
+
+    def test_identity_output_as_multiple_graph_outputs(self):
+        # handle case like this, both Identity nodes are graph outputs,
+        #            Add
+        #           /   \
+        #    Identity   Identity
+        # We at most can remove one Identity for this case.
+        node1 = helper.make_node("Add", ["X", "X"], ["Y"], name="identity")
+        node2 = helper.make_node("Identity", ["Y"], ["Z1"], name="identity2")
+        node3 = helper.make_node("Identity", ["Y"], ["Z2"], name="identity3")
+        graph = helper.make_graph(
+            [node1, node2, node3],
+            "identity-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4, 5))],
+            [helper.make_tensor_value_info("Z1", TensorProto.FLOAT, (2, 3, 4, 5)),
+             helper.make_tensor_value_info("Z2", TensorProto.FLOAT, (2, 3, 4, 5))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_identity_compare(["Z1", "Z2"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                  model_proto, remaining_identity_num=1)
+
+    def test_identity_in_subgraph_non_graph_output(self):
+        node1 = helper.make_node("Add", ["X", "X"], ["Y"], name="add")
+
+        iter_num_value = np.array(1, dtype=np.int64)
+        node2 = helper.make_node(
+            'Constant',
+            inputs=[],
+            outputs=['iterate_num_value'],
+            value=helper.make_tensor(
+                name='iterate_num_value',
+                data_type=TensorProto.INT64,
+                dims=iter_num_value.shape,
+                vals=iter_num_value.flatten().astype(np.int64),
+            ),
+        )
+
+        cond_value = np.array(True, dtype=np.bool)
+        node3 = helper.make_node(
+            'Constant',
+            inputs=[],
+            outputs=['cond_value'],
+            value=helper.make_tensor(
+                name='cond_value',
+                data_type=TensorProto.BOOL,
+                dims=iter_num_value.shape,
+                vals=cond_value.flatten().astype(np.bool),
+            ),
+        )
+
+        # sub graph
+        sub_node1 = helper.make_node("Add", ["loop_var_1", "loop_var_1"], ["SubY"], name="sub_add")
+        sub_node2 = helper.make_node("Identity", ["SubY"], ["SubIdentity1"], name="sub_identity_1")
+        sub_node3 = helper.make_node("Identity", ["SubIdentity1"], ["loop_var_out_1"], name="sub_identity_2")
+        sub_node4 = helper.make_node("Identity", ["loop_condition"], ["loop_cond_output"], name="sub_identity_3")
+        sub_graph = helper.make_graph(
+            [sub_node1, sub_node2, sub_node3, sub_node4],
+            "identity_subgraph-test",
+            [helper.make_tensor_value_info("loop_iter_num", TensorProto.INT64, (1,)),  # iteration_num
+             helper.make_tensor_value_info("loop_condition", TensorProto.BOOL, ()),  # condition
+             helper.make_tensor_value_info("loop_var_1", TensorProto.FLOAT, ()),  # loop-carried dependency
+             ],
+            [helper.make_tensor_value_info("loop_cond_output", TensorProto.BOOL, ()),
+             helper.make_tensor_value_info("loop_var_out_1", TensorProto.FLOAT, ())
+            ],
+        )
+        # sub graph ends
+
+        loop_node = helper.make_node("Loop", ["iterate_num_value", "cond_value", "Y"], ["loop_var_1_output"],
+                                     name="loop", body=sub_graph)
+
+        node4 = helper.make_node("Identity", ["loop_var_1_output"], ["Z"], name="identity")
+        node5 = helper.make_node("Shape", ["Z"], ["Z1"], name="shape")
+
+        graph = helper.make_graph(
+            [node1, node2, node3, loop_node, node4, node5],
+            "identity-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4, 5))],
+            [helper.make_tensor_value_info("Z1", TensorProto.INT64, [4])],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_identity_compare(["Z1"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                  model_proto, remaining_identity_num=0)
 
+    # Tranpose Optimizer Tests End
 
 if __name__ == "__main__":
     unittest_main()

tf2onnx/convert.py

@@ -126,8 +126,8 @@ def main():
 
     model_proto = g.make_model("converted from {}".format(args.input))
 
-    new_model_proto = GraphUtil.opt_transposes_with_graph(g, "converted from {}".format(model_path),
-                                                          optimize=not args.continue_on_error)
+    new_model_proto = GraphUtil.optimize_graph(g, "converted from {}".format(model_path),
+                                               optimize=not args.continue_on_error)
     if new_model_proto:
         model_proto = new_model_proto
     else:

tf2onnx/graph.py

@@ -19,7 +19,7 @@
 from onnx import helper, numpy_helper, optimizer, shape_inference, OperatorSetIdProto, AttributeProto
 from tf2onnx import utils, __version__
 from tf2onnx.utils import port_name, find_opset
-from tf2onnx.optimizer.transpose_optimizer import TransposeOptimizer
+from tf2onnx.optimizer import IdentityOptimizer, TransposeOptimizer
 from tf2onnx.schemas import get_schema
 
 
@@ -854,6 +854,10 @@ def dump_node_statistics(self):
         op_cnt = collections.Counter()
         for n in self.get_nodes():
             op_cnt[n.type] += 1
+            body_graphs = n.get_body_graphs()
+            if body_graphs:
+                for _, b_g in body_graphs.items():
+                    op_cnt += b_g.dump_node_statistics()
 
         return op_cnt
 
@@ -1040,16 +1044,19 @@ class GraphUtil(object):
     """Utilities for Graph manipulation."""
 
     @staticmethod
-    def opt_transposes_with_graph(graph, doc_string, optimize=None, debug=False):
-        """Optimize the graph, eliminating all useless Transpose pairs.
+    def optimize_graph(graph, doc_string, optimize=None, debug=False):
+        """Optimize the graph, for example: eliminating all useless Transpose/Identity pairs.
 
         Returns:
             model proto after optimization, if optimizer run successfully
             or None, if exceptions happen
         """
         try:
-            opt = TransposeOptimizer(graph, output_names=graph.outputs, debug=debug)
-            opt.optimize()
+            opts = [TransposeOptimizer(graph, output_names=graph.outputs, debug=debug),
+                    IdentityOptimizer(graph, output_names=graph.outputs, debug=debug)
+                   ]
+            for opt in opts:
+                opt.optimize()
             model_proto = graph.make_model(doc_string, optimize=optimize)
             return model_proto
         except Exception:
@@ -1060,19 +1067,22 @@ def opt_transposes_with_graph(graph, doc_string, optimize=None, debug=False):
             return None
 
     @staticmethod
-    def opt_transposes_with_model_proto(onnx_model_proto, debug=False):
-        """Optimize the model proto, eliminating all useless Transpose pairs.
+    def optimize_graph_with_model_proto(onnx_model_proto, debug=False):
+        """Optimize the model proto, for example: eliminating all useless Transpose pairs.
 
         Returns:
             model proto after optimization, if optimizer run successfully
             or None, if exceptions happens
         """
         try:
             kwargs = GraphUtil.get_onnx_model_properties(onnx_model_proto)
-
             g = GraphUtil.create_graph_from_onnx_model(onnx_model_proto)
-            opt = TransposeOptimizer(g, output_names=g.outputs, debug=debug)
-            opt.optimize()
+
+            opts = [TransposeOptimizer(g, output_names=g.outputs, debug=debug),
+                    IdentityOptimizer(g, output_names=g.outputs, debug=debug)
+                   ]
+            for opt in opts:
+                opt.optimize()
 
             model_proto = g.make_model(onnx_model_proto.graph.doc_string,
                                        graph_name=onnx_model_proto.graph.name, **kwargs)