add graph optimizer - const_fold_optimizer

Author: zhijxu-MS

tests/test_optimizers.py

@@ -14,7 +14,6 @@
 from backend_test_base import Tf2OnnxBackendTestBase
 from common import unittest_main, group_nodes_by_type
 
-
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
 
 class OptimizerTests(Tf2OnnxBackendTestBase):
@@ -423,5 +422,69 @@ def test_duplicated_need_multiple_run(self):
                                                 op_type="Log", remaining_op_num=3)
     # Merge Duplicated Nodes Optimizer Tests End
 
+    # Const Fold Optimizer Tests Start
+
+    def test_const_fold_trans_with_const1(self):
+        shape = (6, 6)
+        const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
+                                          vals=np.random.randn(*shape).flatten().astype(np.float32))
+        node1 = helper.make_node("Constant", [], ["const"], value=const_tensor)
+        node2 = helper.make_node("Transpose", ["const"], ["value1"])
+        node3 = helper.make_node("Add", ["value1", "X"], ["res"])
+
+        graph = helper.make_graph(
+            [node1, node2, node3],
+            "test_const_fold_trans_with_const1",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, shape)],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"X": np.random.randn(*shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
+    def test_const_fold_trans_with_const2(self):
+        # need multiple optimization run
+        shape = (6, 6)
+        const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
+                                          vals=np.random.randn(*shape).flatten().astype(np.float32))
+        node1 = helper.make_node("Constant", [], ["const"], value=const_tensor)
+        node2 = helper.make_node("Transpose", ["const"], ["value1"])
+        node3 = helper.make_node("Transpose", ["value1"], ["value2"])
+        node4 = helper.make_node("Add", ["value2", "X"], ["res"])
+
+        graph = helper.make_graph(
+            [node1, node2, node3, node4],
+            "test_const_fold_trans_with_const2",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, shape)],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"X": np.random.randn(*shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
+    def test_const_fold_node_is_output(self):
+        # need multiple optimization run
+        shape = (6, 6)
+        const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
+                                          vals=np.random.randn(*shape).flatten().astype(np.float32))
+        node1 = helper.make_node("Constant", [], ["const"], value=const_tensor)
+        node2 = helper.make_node("Transpose", ["const"], ["value1"])
+        node3 = helper.make_node("Transpose", ["value1"], ["res"])
+
+        graph = helper.make_graph(
+            [node1, node2, node3],
+            "test_const_fold_node_is_output",
+            [],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, shape)],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {},
+                                   model_proto, remaining_transpose_num=0)
+    # Const Fold Optimizer Tests End
+
+
 if __name__ == "__main__":
     unittest_main()

tf2onnx/optimizer/__init__.py

@@ -10,15 +10,18 @@
 import traceback
 from collections import OrderedDict
 
+from tf2onnx.optimizer.const_fold_optimizer import ConstFoldOptimizer
 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),
+    ("fold_const", ConstFoldOptimizer),
     # 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),

tf2onnx/optimizer/const_fold_optimizer.py

@@ -0,0 +1,96 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""const fold Optimizer.
+   if op's inputs are all const then do op computation when building the graph to improve performance
+   for example, input of transpose node is const then we can do transpose statically instead of at runtime
+"""
+
+import logging
+
+from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
+from tf2onnx import utils
+
+# pylint: disable=logging-not-lazy,unused-argument,missing-docstring
+
+_func_map = {}
+
+
+def _register_func(onnx_op):
+    def _internal_fun(func):
+        _func_map[onnx_op] = func
+        return func
+    return _internal_fun
+
+
+class ConstFoldOptimizer(GraphOptimizerBase):
+
+    def __init__(self, debug=False):
+        super(ConstFoldOptimizer, self).__init__("ConstFoldOptimizer", debug)
+        self._log = logging.getLogger("tf2onnx.optimizer.%s" % self._name)
+
+    def _optimize(self, graph):
+        return self._apply_optimization(graph, self._optimize_at_current_graph_level)
+
+    def _optimize_at_current_graph_level(self, graph):
+        graph_changed = True
+        while graph_changed:
+            graph_changed = False
+            ops = graph.get_nodes()
+            for op in ops:
+                if self._can_skip(op):
+                    continue
+                if self._fold_node(op, graph):
+                    graph_changed = True
+        return graph
+
+    @staticmethod
+    def _can_skip(node):
+        if node.is_const() or node.is_graph_input():
+            return True
+
+        skip_type = ["Identity"]
+        if node.type in skip_type:
+            return True
+
+        return False
+
+    def _fold_node(self, node, graph):
+        """ if node's input are all const and it's not graph's output then it can be fold.
+            if node can be fold True will be return indicating that graph is changed
+        """
+        if self._all_inputs_are_const(node.inputs) and not self._is_graph_output(node, graph):
+            process_func = _func_map.get(node.type, self._try_fold)
+            return process_func(node, graph)
+
+        return False
+
+    @staticmethod
+    def _all_inputs_are_const(nodes):
+        return all(node.is_const() for node in nodes if node)
+
+    @staticmethod
+    def _is_graph_output(node, graph):
+        node_out_set = set(node.output)
+        graph_out_set = set(graph.outputs)
+        return node_out_set.intersection(graph_out_set)
+
+    @staticmethod
+    def _replace_node_with_const(node, graph, val):
+        const_node = graph.make_const(utils.make_name("const_fold_opt"), val)
+        graph.replace_all_inputs(graph.get_nodes(), node.output[0], const_node.output[0])
+        graph.remove_node(node.name)
+
+    @staticmethod
+    @_register_func(onnx_op="Transpose")
+    def _fold_transpose(node, graph):
+        const_val = node.inputs[0].get_tensor_value(as_list=False)
+        perm_attr = node.get_attr("perm")
+        perm = perm_attr.ints if perm_attr else None
+        const_val_after_trans = const_val.transpose(perm)
+        ConstFoldOptimizer._replace_node_with_const(node, graph, const_val_after_trans)
+        return True
+
+    def _try_fold(self, node, graph):
+        self._log.warning("need to add function to fold op %s whose op_type is %s", node.name, node.type)
+        return False

tf2onnx/optimizer/optimizer_base.py

@@ -17,6 +17,8 @@ def __init__(self, name, debug=False):
     def optimize(self, graph):
         original_node_statistics = graph.dump_node_statistics()
         graph = self._optimize(graph)
+        graph.delete_unused_nodes(graph.outputs)
+        graph.topological_sort(graph.get_nodes())
         node_statistics = graph.dump_node_statistics()
         self._print_stat_diff(original_node_statistics, node_statistics)
         return graph