refactor

Author: zhijxu-MS

tests/backend_test_base.py

@@ -18,7 +18,7 @@
 from common import get_test_config
 from tf2onnx import utils
 from tf2onnx.tfonnx import process_tf_graph, tf_optimize
-from tf2onnx.graph import GraphUtil
+from tf2onnx import optimizer
 
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
@@ -141,7 +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 = GraphUtil.optimize_graph(g)
+            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, debug=debug).make_model("optimized")
+                new_model_proto = optimizer.optimize_graph(onnx_graph, debug=debug).make_model("optimized")
                 if new_model_proto:
                     model_proto = new_model_proto
                 else:

tf2onnx/optimizer/identity_optimizer.py

@@ -8,18 +8,22 @@
 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,unused-variable
 
 
-class IdentityOptimizer(object):
+class IdentityOptimizer(GraphOptimizerBase):
     """Identity Optimizer."""
 
     def __init__(self, debug=False):
-        self._debug = debug
+        super(IdentityOptimizer, self).__init__(debug)
+        self._name = "IdentityOptimizer"
+
         self._g = None
 
     def optimize(self, graph):

tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py

@@ -8,6 +8,7 @@
 """
 
 from collections import defaultdict, namedtuple
+import logging
 
 from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
 
@@ -18,11 +19,14 @@
 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)
-
+    _KeyToGroupNodes = namedtuple("key", "type input")
+
+    def __init__(self, debug=False):
+        super(MergeDuplicatedNodesOptimizer, self).__init__(debug)
+        # optimizer should have name and log property
+        self._name = "MergeDuplicatedNodesOptimizer"
+        self._log = logging.getLogger("tf2onnx.optimizer.%s" % self._name)
+        # used internally
         self._graph_can_be_optimized = True
 
     def _optimize(self, graph):
@@ -46,7 +50,7 @@ def _merge_duplicated_nodes(self, 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)
+            res[self._KeyToGroupNodes(node.type, tuple(node.input))].append(node)
         return res
 
     def _del_nodes_if_duplicated(self, nodes_group, graph):

tf2onnx/optimizer/optimizer_base.py

@@ -5,20 +5,13 @@
 
 from __future__ import unicode_literals
 
-import logging
-
-from tf2onnx import utils
-
 
 class GraphOptimizerBase(object):
     """optimizer graph to improve performance
     """
 
-    def __init__(self, name="GraphOptimizerBase", debug=False):
+    def __init__(self, debug=False):
         self._debug = debug
-        utils.make_sure(name is not None, "name of optimizer is needed.")
-        self._name = name
-        self._log = logging.getLogger("tf2onnx.optimizer.%s" % self._name)
 
     def optimize(self, graph):
         original_node_statistics = graph.dump_node_statistics()

tf2onnx/optimizer/transpose_optimizer.py

@@ -11,6 +11,7 @@
 import numpy as np
 
 from tf2onnx import utils
+from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
 
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger("tf2onnx.optimizer.transpose_optimizer")
@@ -35,11 +36,12 @@ def is_useless_transpose(transpose_node):
     return transpose_node.type == "Transpose" and perm_attr and perm_attr.ints == list(range(len(perm_attr.ints)))
 
 
-class TransposeOptimizer(object):
+class TransposeOptimizer(GraphOptimizerBase):
     """Transpose Optimizer."""
 
     def __init__(self, debug=False):
-        self._debug = debug
+        super(TransposeOptimizer, self).__init__(debug)
+        self._name = "TransposeOptimizer"
         self._handler_map = {}
         self._force_stop = {}