Merge pull request #584 from zhijxu-MS/add_loop_optimizer

enhance optimizer

Author: zhijxu-MS

tests/test_optimizers.py

@@ -50,6 +50,21 @@ def run_and_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
             self.assertEqual(expected_val.shape, actual_val.shape)
 
         return new_proto
+
+    @staticmethod
+    def _make_onnx_const(np_val, output_name):
+        node = helper.make_node(
+            'Constant',
+            inputs=[],
+            outputs=[output_name],
+            value=helper.make_tensor(
+                name=output_name,
+                data_type=utils.map_numpy_to_onnx_dtype(np_val.dtype),
+                dims=np_val.shape,
+                vals=np_val.flatten().astype(np_val.dtype),
+            ),
+        )
+        return node
     # Tranpose Optimizer Tests Start
 
     def run_transpose_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
@@ -304,6 +319,55 @@ def test_transpose_with_squeeze4(self):
         self.run_transpose_compare(["Z"], {"X": np.random.randn(3, 1, 1, 5).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    def test_transpose_with_loop(self):
+        def _define_loop_graph(external_inputs):
+            # external_inputs: external node which will be used by this graph
+            # graph without loop carried
+            # computation
+            # for(...){a = external_inputs[i]; b = trans(a), c = squeeze(b)}, c is scan output
+            node1 = helper.make_node("Gather", [external_inputs[0], "loop_iter_num"], ["Y0"])
+            node2 = helper.make_node("Transpose", ["Y0"], ["Z0"], perm=[0, 2, 3, 1])
+            # graph output
+            node3 = helper.make_node("Squeeze", ["Z0"], ["scan_output"], axes=[0])
+            node4 = helper.make_node("Identity", ["loop_condition"], ["loop_cond_output"])
+            node5 = helper.make_node("Identity", ["loop_condition"], ["loop_carried_output"])
+
+            graph = helper.make_graph(
+                [node1, node2, node3, node4, node5],
+                "loop_subgraph",
+                [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_carried", TensorProto.BOOL, ())  # loop_carried
+                 ],
+                [helper.make_tensor_value_info("loop_cond_output", TensorProto.BOOL, ()),
+                 helper.make_tensor_value_info("loop_carried_output", TensorProto.BOOL, ()),
+                 helper.make_tensor_value_info("scan_output", TensorProto.FLOAT, ["unknown"] * 3)
+                 ],
+            )
+            return graph
+
+        def _make_loop(external_inputs, outputs):
+            trip_cnt = self._make_onnx_const(np.array(10, dtype=np.int64), "trip_cnt")
+            cond = self._make_onnx_const(np.array(True, dtype=np.bool), "cond")
+            sub_graph = _define_loop_graph(external_inputs)
+            loop_node = helper.make_node("Loop", ["trip_cnt", "cond", "cond"], outputs,
+                                         name="loop", body=sub_graph)
+            return trip_cnt, cond, loop_node
+
+        nodes = _make_loop(["array"], ["loop_carried", "scan_out"])
+        res = helper.make_node("Transpose", ["scan_out"], ["Y"], perm=[0, 3, 1, 2], name="trans")
+
+        graph = helper.make_graph(
+            [*nodes, res],
+            "transpose_with_loop",
+            [helper.make_tensor_value_info("array", TensorProto.FLOAT, ["unknow"] * 4)],
+            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, ["unknow"] * 4)],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["Y"], {"array": np.random.randn(10, 3, 4, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
     def test_trans_output_as_graph_outputs(self):
         """
         If transpose's output is graph's output, don't optimize it.

tf2onnx/optimizer/__init__.py

@@ -13,12 +13,14 @@
 from .identity_optimizer import IdentityOptimizer
 from .merge_duplicated_nodes_optimizer import MergeDuplicatedNodesOptimizer
 from .transpose_optimizer import TransposeOptimizer
+from .loop_optimizer import LoopOptimizer
 from .. import logging
 
 # optimizer sequence need to be considered carefully
 _optimizers = OrderedDict([
     ("optimize_transpose", TransposeOptimizer),
     ("fold_constants", ConstFoldOptimizer),
+    ("loop_optimizer", LoopOptimizer),
     # merge_duplication should be used after optimize_transpose
     # for optimize_transpose may have some trans nodes that can be merge
     ("merge_duplication", MergeDuplicatedNodesOptimizer),
@@ -37,14 +39,20 @@ def optimize_graph(graph):
 
     before = graph.dump_node_statistics()
     opts = _get_optimizers()
-    for name, factory in opts.items():
-        try:
-            logger.verbose("Apply %s", name)
-            current = copy.deepcopy(graph)
-            graph = factory().optimize(current)
-        except Exception:  # pylint: disable=broad-except
-            # if current optimizer fails, continue with other optimizers
-            logger.warning("Failed to apply %s", name, exc_info=1)
+    continue_flag = True
+    while continue_flag:
+        continue_flag = False
+        for name, factory in opts.items():
+            try:
+                logger.verbose("Apply %s", name)
+                current = copy.deepcopy(graph)
+                opt = factory()
+                graph = opt.optimize(current)
+                continue_flag = continue_flag or opt.graph_been_opt
+
+            except Exception:  # pylint: disable=broad-except
+                # if current optimizer fails, continue with other optimizers
+                logger.warning("Failed to apply %s", name, exc_info=1)
 
     after = graph.dump_node_statistics()
     diff = copy.deepcopy(after)

tf2onnx/optimizer/const_fold_optimizer.py

@@ -42,6 +42,7 @@ def _optimize_at_current_graph_level(self, graph):
                     continue
                 if self._fold_node(op, graph):
                     graph_changed = True
+                    self.graph_been_opt = True
         return graph
 
     @staticmethod

tf2onnx/optimizer/identity_optimizer.py

@@ -39,6 +39,8 @@ def _optimize_at_current_graph_level(self, g):
                 else:
                     ret = self._handle_non_graph_output_identity(g, n)
                 has_update = ret
+                if ret:
+                    self.graph_been_opt = True
         return g
 
     @staticmethod

tf2onnx/optimizer/loop_optimizer.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""Loop Optimizer.
+   some op in loop's body graph can be moved out to the loop
+"""
+
+from tf2onnx.utils import make_name, make_sure
+from .optimizer_base import GraphOptimizerBase
+
+
+# pylint: disable=logging-not-lazy,unused-argument,missing-docstring,unused-variable,arguments-differ
+
+
+class LoopOptimizer(GraphOptimizerBase):
+    """Loop Optimizer."""
+
+    # a lot of terms used here come from loop's onnx spec
+    # https://github.com/onnx/onnx/blob/master/docs/Operators.md#Loop
+    def __init__(self):  # pylint: disable=useless-super-delegation
+        super(LoopOptimizer, self).__init__()
+
+    def _optimize(self, graph):
+        return self._apply_optimization(graph, self._optimize_at_current_graph_level)
+
+    def _optimize_at_current_graph_level(self, g):
+        has_update = True
+        while has_update:
+            has_update = False
+            nodes = [n for n in g.get_nodes() if n.type == "Loop"]
+            for n in nodes:
+                has_update_tmp = self._try_move_transpose_out_of_body_graph(n)
+                if has_update_tmp:
+                    has_update = True
+                    self.graph_been_opt = True
+        return g
+
+    @staticmethod
+    def consumer_nodes_num(graph, node):
+        make_sure(len(node.output) == 1, "only consider node with only one output")
+        res = len(graph.find_output_consumers(node.output[0]))
+        return res
+
+    def _try_move_transpose_out_of_body_graph(self, loop_node):
+        # output node of body graph can be loop-carried-dependent, if so it can't be move out of the body graph
+        # return True if moving some nodes successfully
+        # for now, we only consider moving transpose
+        body_graph = loop_node.get_body_graphs()["body"]
+        parent_graph = loop_node.graph
+        scan_nodes_name_in_body, scan_node_in_parent = self._scan_outputs(loop_node)
+        scan_nodes = [body_graph.get_node_by_output(name) for name in scan_nodes_name_in_body]
+        graph_is_changed = False
+        for node, name_in_parent in zip(scan_nodes, scan_node_in_parent):
+            # 1 delete node in body graph if possible
+            # only consider two case: trans is output, or transpose > identity > output
+            need_process = False
+            if node.type == "Transpose" and self.consumer_nodes_num(body_graph, node) <= 1:
+                trans = node
+                new_output = node.input[0]
+                body_graph.remove_node(node.name)
+                need_process = True
+            elif node.type == "Identity" and node.inputs[0].type == "Transpose" \
+                    and self.consumer_nodes_num(body_graph, node) <= 1\
+                    and self.consumer_nodes_num(body_graph, node.inputs[0]) <= 1:
+                trans = node.inputs[0]
+                new_output = node.inputs[0].input[0]
+                body_graph.remove_node(node.inputs[0].name)
+                body_graph.remove_node(node.name)
+                need_process = True
+
+            if need_process:
+                # 2 correct body graph's output
+                body_outputs = body_graph.outputs
+                body_outputs[body_outputs.index(node.output[0])] = new_output
+                # 3 insert new node in parent graph
+                ori_perm = list(trans.get_attr("perm").ints)
+                new_perm = [0] + [i + 1 for i in ori_perm]  # body output's rank is m > rank of loop's output is m+1
+                name = make_name("trans_moved_from_loop_body")
+                _ = parent_graph.insert_new_node_on_output("Transpose", name_in_parent, name, perm=new_perm)
+                graph_is_changed = True
+
+        return graph_is_changed
+
+    @classmethod
+    def _scan_outputs(cls, loop):
+        # loop has 2+N inputs; loop has N+K outputs;
+        # loop's body graph has 1+N+K outputs
+        loop_carried = len(loop.input) - 2
+        body_graph = loop.get_body_graphs()["body"]
+        return body_graph.outputs[loop_carried + 1:], loop.output[loop_carried:]

tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py

@@ -32,6 +32,8 @@ 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)
+            if self._graph_can_be_optimized:
+                self.graph_been_opt = True
         return graph
 
     def _merge_duplicated_nodes(self, graph):

tf2onnx/optimizer/optimizer_base.py

@@ -16,6 +16,7 @@ class GraphOptimizerBase(object):
 
     def __init__(self):
         self._logger = logging.getLogger('.'.join(__name__.split('.')[:-1] + [self.__class__.__name__]))
+        self._graph_been_opt = False
 
     @property
     def logger(self):
@@ -25,6 +26,14 @@ def logger(self):
     def is_debug_mode(self):
         return utils.is_debug_mode()
 
+    @property
+    def graph_been_opt(self):
+        return self._graph_been_opt
+
+    @graph_been_opt.setter
+    def graph_been_opt(self, value):
+        self._graph_been_opt = value
+
     def optimize(self, graph):
         """ Optimize graph, return optimized graph. """
         before = graph.dump_node_statistics()

tf2onnx/optimizer/transpose_optimizer.py

@@ -147,6 +147,7 @@ def _optimize_at_current_graph_level(self, graph):
                 if is_nhwc_transpose(n):
                     if self._handle_nhwc_tranpose(n):
                         no_action = False
+                        self.graph_been_opt = True
                         iteration_cnt += 1
                         # need break, because handler may change nodes set, making the n stale object
                         # referencing already deleted elements