Merge pull request #396 from zhijxu-MS/push_branch

add graph optimizer "merge duplicated nodes"

Author: nbcsm

tests/backend_test_base.py

@@ -18,6 +18,7 @@
 from common import get_test_config
 from tf2onnx import utils
 from tf2onnx.tfonnx import process_tf_graph, tf_optimize
+from tf2onnx import optimizer
 
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
@@ -140,6 +141,7 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
         with tf.Session() as sess:
             g = process_tf_graph(sess.graph, opset=self.config.opset, output_names=output_names_with_port,
                                  target=self.config.target, **process_args)
+            g = optimizer.optimize_graph(g)
             actual = self._run_backend(g, output_names_with_port, onnx_feed_dict)
 
         for expected_val, actual_val in zip(expected, actual):

tests/run_pretrained_models.py

@@ -30,7 +30,7 @@
 import tf2onnx
 from tf2onnx import loader
 from tf2onnx import utils
-from tf2onnx.graph import GraphUtil
+from tf2onnx import optimizer
 from tf2onnx.tfonnx import process_tf_graph
 
 # pylint: disable=broad-except,logging-not-lazy,unused-argument,unnecessary-lambda
@@ -270,7 +270,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.optimize_graph(onnx_graph, "test", debug=debug)
+                new_model_proto = optimizer.optimize_graph(onnx_graph, debug=debug).make_model("optimized")
                 if new_model_proto:
                     model_proto = new_model_proto
                 else:

tests/test_optimizers.py

@@ -27,7 +27,7 @@ def run_and_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
 
         origin_model_path = self.save_onnx_model(origin_proto, onnx_feed_dict, postfix="_origin")
 
-        new_proto = GraphUtil.optimize_graph_with_model_proto(origin_proto)
+        new_proto = GraphUtil.optimize_model_proto(origin_proto)
 
         self.assertTrue(new_proto, msg="model proto after optimizer should not be None")
 
@@ -287,7 +287,115 @@ def test_identity_in_subgraph_non_graph_output(self):
         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
+    # Identity Optimizer Tests End
+
+    # Merge Duplicated Nodes Optimizer Tests Start
+
+    def run_merge_duplicated_nodes_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
+                                           op_type=None, remaining_op_num=None, debug=False, rtol=1e-07):
+        self.run_and_compare(output_names_with_port, onnx_feed_dict, origin_proto, op_type=op_type,
+                             remaining_op_num=remaining_op_num, debug=debug, rtol=rtol)
+
+    def test_duplicated_duplicated_input(self):
+        # same input or not
+        node0 = helper.make_node('Add', inputs=["X", "X"], outputs=["value0"])
+        node1 = helper.make_node('Add', inputs=["X", "X"], outputs=["value1"])
+        node2 = helper.make_node('Add', inputs=["value1", "X"], outputs=["value2"])
+        node3 = helper.make_node("Mul", ["value0", "value2"], ["value3"])
+        node4 = helper.make_node("Mul", ["value1", "value3"], ["OUT"])
+
+        graph = helper.make_graph(
+            [node0, node1, node2, node3, node4],
+            "test_duplicated_duplicated_input",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (5, 5))],
+            [helper.make_tensor_value_info("OUT", TensorProto.FLOAT, (5, 5))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_merge_duplicated_nodes_compare(["OUT"], {"X": np.random.randn(5, 5).astype(np.float32)}, model_proto,
+                                                op_type="Add", remaining_op_num=2)
+
+    def test_duplicated_duplicated_attributes(self):
+        # same attr or not
+        node0 = helper.make_node('ReduceSum', inputs=["X"], outputs=["value0"], axes=[0], keepdims=0)
+        node1 = helper.make_node('ReduceSum', inputs=["X"], outputs=["value1"], axes=[0], keepdims=0)
+        node2 = helper.make_node('ReduceSum', inputs=["X"], outputs=["value2"], axes=[1], keepdims=0)
+        node3 = helper.make_node('Add', inputs=["value0", "value1"], outputs=["value3"])
+        node4 = helper.make_node("Mul", ["value2", "value3"], ["OUT"])
+
+        graph = helper.make_graph(
+            [node0, node1, node2, node3, node4],
+            "test_duplicated_duplicated_attributes",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (5, 5))],
+            [helper.make_tensor_value_info("OUT", TensorProto.FLOAT, (5,))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_merge_duplicated_nodes_compare(["OUT"], {"X": np.random.randn(5, 5).astype(np.float32)}, model_proto,
+                                                op_type="ReduceSum", remaining_op_num=2)
+
+    def test_duplicated_node_is_graph_output(self):
+        node0 = helper.make_node('Add', inputs=["X", "X"], outputs=["value0"])
+        node1 = helper.make_node('Add', inputs=["X", "X"], outputs=["value1"])
+        node2 = helper.make_node('Add', inputs=["value1", "X"], outputs=["value2"])
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "test_duplicated_node_is_graph_output",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (5, 5))],
+            [helper.make_tensor_value_info("value1", TensorProto.FLOAT, (5, 5)),
+             helper.make_tensor_value_info("value2", TensorProto.FLOAT, (5, 5))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_merge_duplicated_nodes_compare(["value1", "value2"],
+                                                {"X": np.random.randn(5, 5).astype(np.float32)}, model_proto,
+                                                op_type="Add", remaining_op_num=2)
+
+    def test_duplicated_different_output_length(self):
+        node0 = helper.make_node('Dropout', inputs=["X"], outputs=["value0"])
+        node1 = helper.make_node('Dropout', inputs=["X"], outputs=["value1", "mask"])
+        node2 = helper.make_node('Dropout', inputs=["value1"], outputs=["value2"])
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "test_duplicated_different_output_length",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,))],
+            [helper.make_tensor_value_info("value1", TensorProto.FLOAT, (5,)),
+             helper.make_tensor_value_info("mask", TensorProto.BOOL, (5,)),
+             helper.make_tensor_value_info("value2", TensorProto.FLOAT, (5,))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_merge_duplicated_nodes_compare(["value1", "mask", "value2"],
+                                                {"X": np.random.randn(5,).astype(np.float32)},
+                                                model_proto,
+                                                op_type="Dropout", remaining_op_num=2)
+
+    def test_duplicated_need_multiple_run(self):
+        node00 = helper.make_node('Log', inputs=["X"], outputs=["value00"])
+        node01 = helper.make_node('Log', inputs=["value00"], outputs=["value01"])
+        node02 = helper.make_node('Log', inputs=["value01"], outputs=["value02"])
+
+        node10 = helper.make_node('Log', inputs=["X"], outputs=["value10"])
+        node11 = helper.make_node('Log', inputs=["value10"], outputs=["value11"])
+        node12 = helper.make_node('Log', inputs=["value11"], outputs=["value12"])
+
+        res = helper.make_node('Add', inputs=["value02", "value12"], outputs=["res"])
+
+        graph = helper.make_graph(
+            [node00, node01, node02, node10, node11, node12, res],
+            "test_duplicated_node_is_graph_output",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (5,))],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, (5,))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_merge_duplicated_nodes_compare(["res"], {"X": np.random.randn(5,).astype(np.float32)},
+                                                model_proto,
+                                                op_type="Log", remaining_op_num=3)
+    # Merge Duplicated Nodes Optimizer Tests End
+
 
 if __name__ == "__main__":
     unittest_main()

tf2onnx/convert.py

@@ -117,10 +117,9 @@ def main():
                              output_names=outputs,
                              inputs_as_nchw=args.inputs_as_nchw)
 
-    model_proto = g.make_model("converted from {}".format(args.input))
+    model_proto = g.make_model("converted from {}".format(model_path))
 
-    new_model_proto = GraphUtil.optimize_graph(g, "converted from {}".format(model_path),
-                                               optimize=not args.continue_on_error)
+    new_model_proto = GraphUtil.optimize_model_proto(model_proto)
     if new_model_proto:
         model_proto = new_model_proto
     else:

tf2onnx/graph.py

@@ -17,10 +17,10 @@
 import six
 import numpy as np
 
-from onnx import helper, numpy_helper, optimizer, shape_inference, OperatorSetIdProto, AttributeProto
+from onnx import helper, numpy_helper, shape_inference, OperatorSetIdProto, AttributeProto
 from tf2onnx import utils, __version__
 from tf2onnx.utils import port_name, find_opset
-from tf2onnx.optimizer import IdentityOptimizer, TransposeOptimizer
+from tf2onnx import optimizer
 from tf2onnx.schemas import get_schema
 
 logging.basicConfig(level=logging.INFO)
@@ -1042,30 +1042,11 @@ class GraphUtil(object):
     """Utilities for Graph manipulation."""
 
     @staticmethod
-    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:
-            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:
-            # degradation to non-optimized model proto
-            type_, value_, traceback_ = sys.exc_info()
-            ex_ext = traceback.format_exception(type_, value_, traceback_)
-            print("NON-CRITICAL error in optimizer: ", ex_ext)
-            return None
+    def optimize_graph(graph, debug=False):
+        return optimizer.optimize_graph(graph, debug)
 
     @staticmethod
-    def optimize_graph_with_model_proto(onnx_model_proto, debug=False):
+    def optimize_model_proto(onnx_model_proto, debug=False):
         """Optimize the model proto, for example: eliminating all useless Transpose pairs.
 
         Returns:
@@ -1074,16 +1055,10 @@ def optimize_graph_with_model_proto(onnx_model_proto, debug=False):
         """
         try:
             kwargs = GraphUtil.get_onnx_model_properties(onnx_model_proto)
-            g = GraphUtil.create_graph_from_onnx_model(onnx_model_proto)
-
-            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)
+            graph = GraphUtil.create_graph_from_onnx_model(onnx_model_proto)
+            graph = GraphUtil.optimize_graph(graph, debug)
+            model_proto = graph.make_model(onnx_model_proto.graph.doc_string,
+                                           graph_name=onnx_model_proto.graph.name, **kwargs)
 
             if onnx_model_proto.metadata_props:
                 metadata_props = {p.key: p.value for p in onnx_model_proto.metadata_props}
@@ -1123,11 +1098,12 @@ def create_graph_from_onnx_model(onnx_model_proto):
         # apply shape inference on the model
         inferred_model = shape_inference.infer_shapes(onnx_model_proto)
         graph_proto = inferred_model.graph
-        main_graph = GraphUtil.create_graph_from_onnx_graph(graph_proto)
+        opset_version = onnx_model_proto.opset_import[0].version
+        main_graph = GraphUtil.create_graph_from_onnx_graph(graph_proto, opset_version)
         return main_graph
 
     @staticmethod
-    def create_graph_from_onnx_graph(graph_proto):
+    def create_graph_from_onnx_graph(graph_proto, opset_version=None):
         """Create Graph loading onnx graph proto."""
         output_shapes = {}
         output_dtypes = {}
@@ -1154,15 +1130,15 @@ def create_graph_from_onnx_graph(graph_proto):
         for n in graph_proto.output:
             output_names.append(n.name)
 
-        g = Graph(nodes_to_append, output_shapes, output_dtypes, None, None, None, output_names)
+        g = Graph(nodes_to_append, output_shapes, output_dtypes, None, opset_version, None, output_names)
         const_nodes = GraphUtil._parse_graph_initializer(g, graph_proto)
         GraphUtil._parse_graph_input(g, graph_proto, [n.name for n in const_nodes])
 
         for n in g.get_nodes():
             for attr_name, attr_val in n.attr.items():
                 if attr_val.HasField('g'):
                     # it was assumed that the a.g has inferred shapes/dtypes.
-                    sub_g = GraphUtil.create_graph_from_onnx_graph(attr_val.g)
+                    sub_g = GraphUtil.create_graph_from_onnx_graph(attr_val.g, opset_version)
                     n.set_body_graph_as_attr(attr_name, sub_g)
         return g
 

tf2onnx/optimizer/__init__.py

@@ -6,10 +6,39 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-from .identity_optimizer import IdentityOptimizer
-from .transpose_optimizer import TransposeOptimizer
+import sys
+import traceback
+from collections import OrderedDict
 
-__all__ = [
-    "IdentityOptimizer",
-    "TransposeOptimizer",
-]
+from tf2onnx.optimizer.identity_optimizer import IdentityOptimizer
+from tf2onnx.optimizer.merge_duplicated_nodes_optimizer import MergeDuplicatedNodesOptimizer
+from tf2onnx.optimizer.transpose_optimizer import TransposeOptimizer
+
+# pylint: disable=missing-docstring, broad-except
+
+# optimizer sequence need to be considered carefully
+_optimizers = OrderedDict([
+    ("transpose_opt", TransposeOptimizer),
+    # merge_duplicated_nodes should be used after transpose_opt
+    # for transpose_opt may have some trans nodes that can be merge
+    ("merge_duplicated_nodes", MergeDuplicatedNodesOptimizer),
+    ("identity_opt", IdentityOptimizer),
+])
+
+
+def optimize_graph(graph, debug=False):
+    try:
+        opts = _get_optimizers()
+        for opt in opts.values():
+            graph = opt(debug=debug).optimize(graph)
+        return graph
+    except Exception:
+        # degradation to non-optimized model proto
+        type_, value_, traceback_ = sys.exc_info()
+        ex_ext = traceback.format_exception(type_, value_, traceback_)
+        print("NON-CRITICAL error in optimizer: ", ex_ext)
+        return None
+
+
+def _get_optimizers():
+    return _optimizers

tf2onnx/optimizer/identity_optimizer.py

@@ -8,30 +8,32 @@
 from __future__ import unicode_literals
 import logging
 
-logging.basicConfig(level=logging.INFO)
+from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
+
+
 log = logging.getLogger("tf2onnx.optimizer.identity_optimizer")
 
 
-# pylint: disable=logging-not-lazy,unused-argument,missing-docstring
-# FIXME:
-# pylint: disable=unused-variable
+# pylint: disable=logging-not-lazy,unused-argument,missing-docstring,unused-variable,arguments-differ
 
 
-class IdentityOptimizer(object):
+class IdentityOptimizer(GraphOptimizerBase):
     """Identity Optimizer."""
 
-    def __init__(self, main_graph, output_names, debug=False):
-        self._g = main_graph
-        self._output_names = [name.split(":")[0] for name in output_names]
-        self._debug = debug
+    def __init__(self, debug=False):
+        super(IdentityOptimizer, self).__init__("IdentityOptimizer", debug)
 
-    def optimize(self):
+        self._g = None
+
+    def optimize(self, graph):
+        self._g = graph
         previous_counter = self._g.dump_node_statistics()
         self._optimize_recursively(self._g)
         current_counter = self._g.dump_node_statistics()
         identity_cnt = current_counter["Identity"]
         current_counter.subtract(previous_counter)
         log.info(" %d identity op(s) left, ops diff after identity optimization: %s", identity_cnt, current_counter)
+        return self._g
 
     def _optimize_recursively(self, g):
         self._optimize(g)
@@ -64,13 +66,14 @@ def _optimize(self, g):
 
         self._g.topological_sort(self._g.get_nodes())
 
-
-    def _handle_non_graph_output_identity(self, graph, identity):
+    @staticmethod
+    def _handle_non_graph_output_identity(graph, identity):
         graph.replace_all_inputs(graph.get_nodes(), identity.output[0], identity.input[0])
         graph.remove_node(identity.name)
         return True
 
-    def _handle_graph_output_identity(self, graph, identity, graph_outputs):
+    @staticmethod
+    def _handle_graph_output_identity(graph, identity, graph_outputs):
         input_id = identity.input[0]
         input_node = identity.inputs[0]