refactor optimizer

Author: nbcsm

tf2onnx/convert.py

@@ -85,6 +85,8 @@ def main():
     if args.debug:
         utils.set_debug_mode(True)
 
+    logger = logging.getLogger(constants.TF2ONNX_PACKAGE_NAME)
+
     # override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
     # support unknown dimensions.
     utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
@@ -126,11 +128,8 @@ def main():
 
     model_proto = g.make_model("converted from {}".format(model_path))
 
-    new_model_proto = GraphUtil.optimize_model_proto(model_proto)
-    if new_model_proto:
-        model_proto = new_model_proto
-    else:
-        print("NON-CRITICAL, optimizers are not applied successfully")
+    logger.info("\n")
+    model_proto = GraphUtil.optimize_model_proto(model_proto)
 
     # write onnx graph
     if args.output:

tf2onnx/graph.py

@@ -12,8 +12,6 @@
 import collections
 import copy
 import logging
-import sys
-import traceback
 import six
 import numpy as np
 
@@ -1065,27 +1063,18 @@ def optimize_model_proto(onnx_model_proto):
         """Optimize the model proto, for example: eliminating all useless Transpose pairs.
 
         Returns:
-            model proto after optimization, if optimizer run successfully
-            or None, if exceptions happens
+            model proto after optimization
         """
-        try:
-            kwargs = GraphUtil.get_onnx_model_properties(onnx_model_proto)
-            graph = GraphUtil.create_graph_from_onnx_model(onnx_model_proto)
-            graph = GraphUtil.optimize_graph(graph)
-            model_proto = graph.make_model(onnx_model_proto.graph.doc_string,
-                                           graph_name=onnx_model_proto.graph.name, **kwargs)
-
-            if onnx_model_proto.metadata_props:
-                metadata_props = {p.key: p.value for p in onnx_model_proto.metadata_props}
-                helper.set_model_props(model_proto, metadata_props)
-            return model_proto
-        except Exception:
-            # sometimes, onnx shape inference will fail for some reason, in this case,
-            # we just log the error, and skip the transpose optimizer.
-            type_, value_, traceback_ = sys.exc_info()
-            ex_ext = traceback.format_exception(type_, value_, traceback_)
-            print("NON-CRITICAL error in optimizer: ", ex_ext)
-            return None
+        kwargs = GraphUtil.get_onnx_model_properties(onnx_model_proto)
+        graph = GraphUtil.create_graph_from_onnx_model(onnx_model_proto)
+        graph = GraphUtil.optimize_graph(graph)
+        model_proto = graph.make_model(onnx_model_proto.graph.doc_string,
+                                       graph_name=onnx_model_proto.graph.name, **kwargs)
+
+        if onnx_model_proto.metadata_props:
+            metadata_props = {p.key: p.value for p in onnx_model_proto.metadata_props}
+            helper.set_model_props(model_proto, metadata_props)
+        return model_proto
 
     @staticmethod
     def get_onnx_model_properties(onnx_model_proto):

tf2onnx/optimizer/__init__.py

@@ -6,43 +6,51 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import sys
-import traceback
 from collections import OrderedDict
+import copy
 
-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
+from .const_fold_optimizer import ConstFoldOptimizer
+from .identity_optimizer import IdentityOptimizer
+from .merge_duplicated_nodes_optimizer import MergeDuplicatedNodesOptimizer
+from .transpose_optimizer import TransposeOptimizer
+from .. import logging
 
 # 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),
-    ("identity_opt", IdentityOptimizer),
+    ("reduce_transpose", TransposeOptimizer),
+    ("fold_constants", ConstFoldOptimizer),
+    # merge_duplication should be used after reduce_transpose
+    # for reduce_transpose may have some trans nodes that can be merge
+    ("merge_duplication", MergeDuplicatedNodesOptimizer),
+    ("reduce_identity", IdentityOptimizer),
 ])
 
 
-def optimize_graph(graph):
-    try:
-        opts = _get_optimizers()
-        for opt in opts.values():
-            graph = opt().optimize(graph)
-
-        graph.update_proto()
-        return graph
-    except Exception:
-        # degradation to non-optimized model proto
-        type_, value_, traceback_ = sys.exc_info()
-        ex_ext = traceback.format_exception(type_, value_, traceback_)
-        print("NON-CRITICAL error in optimizer: ", ex_ext)
-        return None
-
-
 def _get_optimizers():
     return _optimizers
+
+
+def optimize_graph(graph):
+    """ Optimize graph, return optimized graph. No throw. """
+    logger = logging.getLogger(__name__)
+    logger.info("Optimizing ONNX model")
+
+    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)
+
+    after = graph.dump_node_statistics()
+    diff = copy.deepcopy(after)
+    diff.subtract(before)
+    diff = ["{} {} ({}->{})".format(k, str(v) if v < 0 else '+' + str(v), before.get(k, 0), after.get(k, 0))
+            for k, v in diff.most_common() if v != 0]
+    logger.info("After optimization: %s", ', '.join(diff) if diff else "no change")
+
+    return graph

tf2onnx/optimizer/const_fold_optimizer.py

@@ -6,8 +6,8 @@
    for example, input of transpose node is const then we can do transpose statically instead of at runtime
 """
 
-from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
-from tf2onnx import utils
+from .. import utils
+from .optimizer_base import GraphOptimizerBase
 
 # pylint: disable=logging-not-lazy,unused-argument,missing-docstring
 

tf2onnx/optimizer/identity_optimizer.py

@@ -7,7 +7,7 @@
 
 from __future__ import unicode_literals
 
-from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
+from .optimizer_base import GraphOptimizerBase
 
 
 # pylint: disable=logging-not-lazy,unused-argument,missing-docstring,unused-variable,arguments-differ
@@ -16,39 +16,20 @@
 class IdentityOptimizer(GraphOptimizerBase):
     """Identity Optimizer."""
 
-    def __init__(self):
+    def __init__(self):  # pylint: disable=useless-super-delegation
         super(IdentityOptimizer, self).__init__()
-        self._g = None
-
-    def optimize(self, graph):
-        self._g = graph
-        previous_counter = self._g.dump_node_statistics()
-        self._optimize_recursively(self._g)
-        current_counter = self._g.dump_node_statistics()
-        identity_cnt = current_counter["Identity"]
-        self.logger.info(" %d identity op(s) left", identity_cnt)
-        self._print_stat_diff(previous_counter, current_counter)
-        return self._g
-
-    def _optimize_recursively(self, g):
-        self._optimize(g)
-        nodes = [n for n in g.get_nodes()]
-        for n in nodes:
-            body_graphs = n.get_body_graphs()
-            if body_graphs:
-                for attr, b_g in body_graphs.items():
-                    self.logger.debug("start handling subgraph of %s's attribute %s", n.name, attr)
-                    self._optimize_recursively(b_g)
-                    self.logger.debug("finish handling subgraph of %s's attribute %s", n.name, attr)
-
-    def _optimize(self, g):
+
+    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 == "Identity"]
             for n in nodes:
                 if n.graph is None:
-                    self.logger.info("node has been removed from this graph, skip")
+                    self.logger.debug("node has been removed from this graph, skip")
                     continue
 
                 graph_outputs = set(n.output).intersection(g.outputs)
@@ -58,8 +39,7 @@ def _optimize(self, g):
                 else:
                     ret = self._handle_non_graph_output_identity(g, n)
                 has_update = ret
-
-        self._g.topological_sort(self._g.get_nodes())
+        return g
 
     @staticmethod
     def _handle_non_graph_output_identity(graph, identity):

tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py

@@ -9,7 +9,7 @@
 
 from collections import defaultdict, namedtuple
 
-from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
+from .optimizer_base import GraphOptimizerBase
 
 # pylint: disable=logging-not-lazy,unused-argument,missing-docstring
 

tf2onnx/optimizer/optimizer_base.py

@@ -5,7 +5,9 @@
 
 from __future__ import unicode_literals
 
-from tf2onnx import logging, utils
+import copy
+
+from .. import logging, utils
 
 
 class GraphOptimizerBase(object):
@@ -24,11 +26,15 @@ def is_debug_mode(self):
         return utils.is_debug_mode()
 
     def optimize(self, graph):
-        original_node_statistics = graph.dump_node_statistics()
+        """ Optimize graph, return optimized graph """
+        before = graph.dump_node_statistics()
+
         graph = self._optimize(graph)
+        graph.update_proto()
         graph.delete_unused_nodes(graph.outputs)
-        node_statistics = graph.dump_node_statistics()
-        self._print_stat_diff(original_node_statistics, node_statistics)
+
+        after = graph.dump_node_statistics()
+        self._print_stat_diff(before, after)
         return graph
 
     def _optimize(self, graph):
@@ -48,14 +54,13 @@ def _apply_optimization(graph, optimize_func):
             body_graphs = node.get_body_graphs()
             if body_graphs:
                 for attr, b_g in body_graphs.items():
-                    b_g = optimize_func(b_g)
+                    b_g = GraphOptimizerBase._apply_optimization(b_g, optimize_func)
                     node.set_body_graph_as_attr(attr, b_g)
         return graph
 
-    def _print_stat_diff(self, nodes_original, nodes_after_optimized):
-        nodes_after_optimized.subtract(nodes_original)
-        res = {}
-        for key, value in nodes_after_optimized.items():
-            if value != 0:
-                res[key] = value
-        self.logger.info("the optimization gain is %s", res)
+    def _print_stat_diff(self, before, after):
+        diff = copy.deepcopy(after)
+        diff.subtract(before)
+        diff = ["{} {} ({}->{})".format(k, str(v) if v < 0 else '+' + str(v), before.get(k, 0), after.get(k, 0))
+                for k, v in diff.most_common() if v != 0]
+        self.logger.verbose(', '.join(diff) if diff else "no change")

tf2onnx/optimizer/transpose_optimizer.py

@@ -8,8 +8,8 @@
 
 import numpy as np
 
-from tf2onnx import utils
-from tf2onnx.optimizer.optimizer_base import GraphOptimizerBase
+from .. import utils
+from .optimizer_base import GraphOptimizerBase
 
 
 # pylint: disable=logging-not-lazy,unused-argument,missing-docstring,abstract-method
@@ -164,10 +164,7 @@ def optimize(self, graph):
 
         current_counter = self._g.dump_node_statistics()
         transpose_cnt = current_counter["Transpose"]
-        self.logger.info(" %d transpose op(s) left", transpose_cnt)
         self._print_stat_diff(previous_counter, current_counter)
-        if transpose_cnt > 2:
-            self.logger.warning("please try add --fold_const to help remove more transpose")
         return self._g
 
     def _initialize_handlers(self):