refactor

Author: zhijxu-MS

tf2onnx/optimizer/identity_optimizer.py

@@ -21,8 +21,7 @@ class IdentityOptimizer(GraphOptimizerBase):
     """Identity Optimizer."""
 
     def __init__(self, debug=False):
-        super(IdentityOptimizer, self).__init__(debug)
-        self._name = "IdentityOptimizer"
+        super(IdentityOptimizer, self).__init__("IdentityOptimizer", debug)
 
         self._g = None
 

tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py

@@ -17,13 +17,12 @@
 
 _KeyToGroupNodes = namedtuple("key", "type input")
 
+
 class MergeDuplicatedNodesOptimizer(GraphOptimizerBase):
     """Remove duplicate nodes.
     """
     def __init__(self, debug=False):
-        super(MergeDuplicatedNodesOptimizer, self).__init__(debug)
-        # optimizer should have name and log property
-        self._name = "MergeDuplicatedNodesOptimizer"
+        super(MergeDuplicatedNodesOptimizer, self).__init__("MergeDuplicatedNodesOptimizer", debug)
         self._log = logging.getLogger("tf2onnx.optimizer.%s" % self._name)
         # used internally
         self._graph_can_be_optimized = True

tf2onnx/optimizer/optimizer_base.py

@@ -10,8 +10,9 @@ class GraphOptimizerBase(object):
     """optimizer graph to improve performance
     """
 
-    def __init__(self, debug=False):
+    def __init__(self, name, debug=False):
         self._debug = debug
+        self._name = name
 
     def optimize(self, graph):
         original_node_statistics = graph.dump_node_statistics()
@@ -23,6 +24,10 @@ def optimize(self, graph):
     def _optimize(self, graph):
         raise NotImplementedError
 
+    @property
+    def name(self):
+        return self._name
+
     @staticmethod
     def _apply_optimization(graph, optimize_func):
         """

tf2onnx/optimizer/transpose_optimizer.py

@@ -40,8 +40,7 @@ class TransposeOptimizer(GraphOptimizerBase):
     """Transpose Optimizer."""
 
     def __init__(self, debug=False):
-        super(TransposeOptimizer, self).__init__(debug)
-        self._name = "TransposeOptimizer"
+        super(TransposeOptimizer, self).__init__("TransposeOptimizer", debug)
         self._handler_map = {}
         self._force_stop = {}