add graph optimizer - merge duplicated nodes

zhijxu-MS · zhijxu-MS · commit 444923fdca4d · 2019-03-27T15:38:37.000+08:00
diff --git a/tests/run_pretrained_models.py b/tests/run_pretrained_models.py
@@ -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 = GraphUtil.optimize_graph(onnx_graph, debug=debug).make_model("optimized")
                 if new_model_proto:
                     model_proto = new_model_proto
                 else:
diff --git a/tests/test_optimizers.py b/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,54 @@ 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],
+            "transpose-merge-test",
+            [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],
+            "transpose-merge-test",
+            [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)
+    # Merge Duplicated Nodes Optimizer Tests End
+
 
 if __name__ == "__main__":
     unittest_main()
diff --git a/tf2onnx/convert.py b/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:
diff --git a/tf2onnx/graph.py b/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,7 +1130,7 @@ 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])
 
diff --git a/tf2onnx/optimizer/__init__.py b/tf2onnx/optimizer/__init__.py
@@ -6,10 +6,38 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-from .identity_optimizer import IdentityOptimizer
-from .transpose_optimizer import TransposeOptimizer
+import sys
+import traceback
 
-__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 = {
+    "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
diff --git a/tf2onnx/optimizer/identity_optimizer.py b/tf2onnx/optimizer/identity_optimizer.py
@@ -12,26 +12,25 @@
 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
 
 
 class IdentityOptimizer(object):
     """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]
+    def __init__(self, debug=False):
         self._debug = debug
+        self._g = None
 
-    def optimize(self):
+    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 +63,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]
 
diff --git a/tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py b/tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py
@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""Merge Duplicated Nodes Optimizer.
+   Remove duplicate nodes except identity nodes which should be handled by identity optimizer.
+   for example, node a is input of node b and node c, and computation of node b, c are same such as "abs" op.
+   then b and c can be merged into one node to avoid duplicated computation
+"""
+
+from collections import defaultdict, namedtuple
+
+from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
+
+
+# pylint: disable=logging-not-lazy,unused-argument,missing-docstring
+
+
+class MergeDuplicatedNodesOptimizer(GraphOptimizerBase):
+    """Remove duplicate nodes.
+    """
+    _key_to_group_nodes = namedtuple("key", "type input")
+
+    def __init__(self, name="MergeDuplicatedNodesOptimizer", debug=False):
+        super(MergeDuplicatedNodesOptimizer, self).__init__(name=name)
+
+        self._graph_can_be_optimized = True
+
+    def _optimize(self, graph):
+        return self._apply_optimization(graph, self._optimize_at_current_graph_level)
+
+    def _optimize_at_current_graph_level(self, graph):
+        while self._graph_can_be_optimized:
+            self._graph_can_be_optimized = False
+            self._merge_duplicated_nodes(graph)
+        return graph
+
+    def _merge_duplicated_nodes(self, graph):
+        # "duplicated" means: op_type, input and attribute are same
+        # while attr is un-hashable so doesn't include it when grouping nodes
+        nodes_groups = self._group_nodes_by_type_inputs(graph)
+        for _, nodes_group in nodes_groups.items():
+            if self._skip_node_type(nodes_group[0]):
+                continue
+            self._del_nodes_if_duplicated(nodes_group, graph)
+
+    def _group_nodes_by_type_inputs(self, graph):
+        res = defaultdict(list)
+        for node in graph.get_nodes():
+            res[self._key_to_group_nodes(node.type, tuple(node.input))].append(node)
+        return res
+
+    def _del_nodes_if_duplicated(self, nodes_group, graph):
+        # input and op type of nodes in same group are same,
+        # and if their attributes are also same then they are duplicated
+        while len(nodes_group) > 1:
+            unprocessed_node = []
+            nodes_to_process = [nodes_group[0]]
+            for node in nodes_group[1:]:
+                if node.attr == nodes_to_process[0].attr:
+                    nodes_to_process.append(node)
+                else:
+                    unprocessed_node.append(node)
+
+            self._merge_nodes_that_are_duplicated(nodes_to_process, graph)
+            nodes_group = unprocessed_node
+
+    def _merge_nodes_that_are_duplicated(self, nodes_to_process, graph):
+        # node's output may not all be used, so have to select the one that uses most of node's outputs
+        nodes_to_process.sort(key=self._len_of_node_output, reverse=True)
+        node_to_retain = nodes_to_process[0]
+        for node_to_delete in nodes_to_process[1:]:
+            # if one of the output is graph's output then it can't be deleted
+            if set(node_to_delete.output).intersection(set(graph.outputs)):
+                continue
+            for old_input, new_input in zip(node_to_delete.output, node_to_retain.output):
+                graph.replace_all_inputs(graph.get_nodes(), old_input, new_input)
+            graph.remove_node(node_to_delete.name)
+            self._graph_can_be_optimized = True
+
+    @staticmethod
+    def _skip_node_type(node):
+        # identity node will be handled by identity optimizer so skip it
+        if node.type in ["Identity"]:
+            return True
+        if node.is_graph_input():
+            return True
+        return False
+
+    @staticmethod
+    def _len_of_node_output(node):
+        return len(node.output)
diff --git a/tf2onnx/optimizer/optimizer_base.py b/tf2onnx/optimizer/optimizer_base.py
diff --git a/tf2onnx/optimizer/transpose_optimizer.py b/tf2onnx/optimizer/transpose_optimizer.py
