refactor code according to review comments

Author: zhijxu-MS

tf2onnx/optimizer/const_fold_optimizer.py

@@ -6,8 +6,6 @@
    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
 
@@ -16,9 +14,9 @@
 _func_map = {}
 
 
-def _register_func(onnx_op):
+def _register_func(op_type):
     def _internal_fun(func):
-        _func_map[onnx_op] = func
+        _func_map[op_type] = func
         return func
     return _internal_fun
 
@@ -27,7 +25,6 @@ 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)
@@ -38,14 +35,18 @@ def _optimize_at_current_graph_level(self, graph):
             graph_changed = False
             ops = graph.get_nodes()
             for op in ops:
-                if self._can_skip(op):
+                if self._should_skip(op):
                     continue
                 if self._fold_node(op, graph):
                     graph_changed = True
         return graph
 
     @staticmethod
-    def _can_skip(node):
+    def _should_skip(node):
+        # only support onnx official op for now, other op such as contrib op not supported for now
+        if not utils.is_onnx_domain(node.domain):
+            return True
+
         if node.is_const() or node.is_graph_input():
             return True
 
@@ -60,9 +61,13 @@ def _fold_node(self, node, graph):
             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)
-
+            process_func = _func_map.get(node.type, None)
+            if process_func:
+                const_val_after_trans = process_func(node, graph)
+                self._replace_node_with_const(node, graph, const_val_after_trans)
+                return True
+            else:
+                self.log.warning("need to add function to fold op %s whose op_type is %s", node.name, node.type)
         return False
 
     @staticmethod
@@ -76,21 +81,20 @@ def _is_graph_output(node, graph):
         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])
+    def _replace_node_with_const(node, graph, vals):
+        utils.make_sure(len(node.output) == len(vals), "length of node outputs and const vals should be same")
+        for old_input, val in zip(node.output, vals):
+            const_node = graph.make_const(utils.make_name("const_fold_opt"), val)
+            graph.set_dtype(const_node.output[0], utils.map_numpy_to_onnx_dtype(val.dtype))
+            graph.set_shape(const_node.output[0], val.shape)
+            graph.replace_all_inputs(graph.get_nodes(), old_input, const_node.output[0])
         graph.remove_node(node.name)
 
     @staticmethod
-    @_register_func(onnx_op="Transpose")
-    def _fold_transpose(node, graph):
+    @_register_func("Transpose")
+    def _fold_transpose(node, graph) -> list:
         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
+        return [const_val_after_trans]

tf2onnx/optimizer/identity_optimizer.py

@@ -6,7 +6,6 @@
 """
 
 from __future__ import unicode_literals
-import logging
 
 from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
 
@@ -18,7 +17,6 @@ class IdentityOptimizer(GraphOptimizerBase):
 
     def __init__(self, debug=False):
         super(IdentityOptimizer, self).__init__("IdentityOptimizer", debug)
-        self._log = logging.getLogger("tf2onnx.optimizer.%s" % self._name)
         self._g = None
 
     def optimize(self, graph):
@@ -27,7 +25,7 @@ def optimize(self, graph):
         self._optimize_recursively(self._g)
         current_counter = self._g.dump_node_statistics()
         identity_cnt = current_counter["Identity"]
-        self._log.info(" %d identity op(s) left", identity_cnt)
+        self.log.info(" %d identity op(s) left", identity_cnt)
         self._print_stat_diff(previous_counter, current_counter)
         return self._g
 
@@ -38,9 +36,9 @@ def _optimize_recursively(self, g):
             body_graphs = n.get_body_graphs()
             if body_graphs:
                 for attr, b_g in body_graphs.items():
-                    self._log.debug("start handling subgraph of %s's attribute %s", n.name, attr)
+                    self.log.debug("start handling subgraph of %s's attribute %s", n.name, attr)
                     self._optimize_recursively(b_g)
-                    self._log.debug("finish handling subgraph of %s's attribute %s", n.name, attr)
+                    self.log.debug("finish handling subgraph of %s's attribute %s", n.name, attr)
 
     def _optimize(self, g):
         has_update = True
@@ -49,7 +47,7 @@ def _optimize(self, g):
             nodes = [n for n in g.get_nodes() if n.type == "Identity"]
             for n in nodes:
                 if n.graph is None:
-                    self._log.info("node has been removed from this graph, skip")
+                    self.log.info("node has been removed from this graph, skip")
                     continue
 
                 graph_outputs = set(n.output).intersection(g.outputs)
@@ -74,12 +72,12 @@ def _handle_graph_output_identity(self, graph, identity, graph_outputs):
 
         if input_node.graph != graph:
             # If input node is in parent graph, we don't handle it now
-            self._log.debug("input node in parent graph, skip")
+            self.log.debug("input node in parent graph, skip")
             return False
 
         if input_node.is_graph_input():
             # Identity between input and output should not be removed.
-            self._log.debug("skip identity between input and output")
+            self.log.debug("skip identity between input and output")
             return False
 
         output_id = identity.output[0]
@@ -88,7 +86,7 @@ def _handle_graph_output_identity(self, graph, identity, graph_outputs):
         if input_id in graph.outputs:
             # input id already be graph output, so we cannot make that be another graph output.
             # this Identity must be kept.
-            self._log.debug("identity input already be graph output")
+            self.log.debug("identity input already be graph output")
             return False
 
         graph.remove_node(identity.name)

tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py

@@ -8,7 +8,6 @@
 """
 
 from collections import defaultdict, namedtuple
-import logging
 
 from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
 
@@ -23,7 +22,6 @@ class MergeDuplicatedNodesOptimizer(GraphOptimizerBase):
     """
     def __init__(self, debug=False):
         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

@@ -4,6 +4,7 @@
 """Graph Optimizer Base"""
 
 from __future__ import unicode_literals
+import logging
 
 
 class GraphOptimizerBase(object):
@@ -13,6 +14,7 @@ class GraphOptimizerBase(object):
     def __init__(self, name, debug=False):
         self._debug = debug
         self._name = name
+        self._log = logging.getLogger("tf2onnx.optimizer.%s" % self._name)
 
     def optimize(self, graph):
         original_node_statistics = graph.dump_node_statistics()
@@ -30,6 +32,10 @@ def _optimize(self, graph):
     def name(self):
         return self._name
 
+    @property
+    def log(self):
+        return self._log
+
     @staticmethod
     def _apply_optimization(graph, optimize_func):
         """
@@ -54,4 +60,4 @@ def _print_stat_diff(self, nodes_original, nodes_after_optimized):
         for key, value in nodes_after_optimized.items():
             if value != 0:
                 res[key] = value
-        self._log.info("the optimization gain is %s", res)
+        self.log.info("the optimization gain is %s", res)

tf2onnx/optimizer/transpose_optimizer.py

@@ -6,7 +6,6 @@
 from __future__ import unicode_literals
 from collections import defaultdict
 
-import logging
 
 import numpy as np
 
@@ -38,7 +37,7 @@ class TransposeOptimizer(GraphOptimizerBase):
 
     def __init__(self, debug=False):
         super(TransposeOptimizer, self).__init__("TransposeOptimizer", debug)
-        self._log = logging.getLogger("tf2onnx.optimizer.%s" % self._name)
+
         self._handler_map = {}
         self._force_stop = {}
 
@@ -156,17 +155,17 @@ def optimize(self, graph):
             if "stop" in self._force_stop and self._force_stop["stop"] == 1:
                 break
 
-        self._log.debug("finish after " + str(iteration_cnt) + " iteration(s)")
+        self.log.debug("finish after " + str(iteration_cnt) + " iteration(s)")
 
         self.merge_duplicated_transposes()
         self.post_optimize_action()
 
         current_counter = self._g.dump_node_statistics()
         transpose_cnt = current_counter["Transpose"]
-        self._log.info(" %d transpose op(s) left", transpose_cnt)
+        self.log.info(" %d transpose op(s) left", transpose_cnt)
         self._print_stat_diff(previous_counter, current_counter)
         if transpose_cnt > 2:
-            self._log.warning("please try add --fold_const to help remove more transpose")
+            self.log.warning("please try add --fold_const to help remove more transpose")
         return self._g
 
     def _initialize_handlers(self):
@@ -217,7 +216,7 @@ def _handle_node_having_branches(self, node):
                 self._g.remove_node(n.name)
             return True
 
-        self._log.debug("input transpose does not have single consumer, skipping...")
+        self.log.debug("input transpose does not have single consumer, skipping...")
         return False
 
     # get the input index of transpose op in node's inputs.
@@ -255,13 +254,13 @@ def _switch_transpose_and_node(self, node, trans):
     # otherwise, it means that we skip handling since it is not in our support set
     def _handle_nhwc_tranpose(self, trans):
         if trans.output[0] in self._g.outputs:
-            log.debug("%s connects to graph outputs, skip", trans.output[0])
+            self.log.debug("%s connects to graph outputs, skip", trans.output[0])
             return False
         out_nodes = self._g.find_output_consumers(trans.output[0])
         if len(out_nodes) == 1:
             p = out_nodes[0]
             if p.name in self._output_names:
-                self._log.debug("cannot move transpose down since it met output node %s", p.name)
+                self.log.debug("cannot move transpose down since it met output node %s", p.name)
                 return False
 
             if p.type in self._handler_map:

tf2onnx/utils.py

@@ -428,6 +428,14 @@ def get_onnx_version():
     return onnx.__version__
 
 
+
 def make_opsetid(domain, version):
     make_sure(isinstance(version, int), "version must be an integer")
     return helper.make_opsetid(domain, version)
+
+
+def is_onnx_domain(domain):
+    if domain is None or domain == "":
+        return True
+    return False
+