Merge pull request #459 from nbcsm/logging

refactor optimizer & refine logging

Author: nbcsm

tests/run_pretrained_models.py

@@ -13,7 +13,6 @@
 import sys
 import tarfile
 import time
-import traceback
 import zipfile
 
 import PIL.Image
@@ -186,7 +185,7 @@ def run_onnxruntime(self, name, model_proto, inputs):
         """Run test against msrt-next backend."""
         import onnxruntime as rt
         model_path = utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto, include_test_data=True)
-        print("\t\t" + model_path)
+        logger.info("Model saved to %s", model_path)
         m = rt.InferenceSession(model_path)
         results = m.run(self.output_names, inputs)
         if self.perf:
@@ -201,12 +200,13 @@ def create_onnx_file(name, model_proto, inputs, outdir):
         os.makedirs(outdir, exist_ok=True)
         model_path = os.path.join(outdir, name + ".onnx")
         utils.save_protobuf(model_path, model_proto)
-        print("\tcreated", model_path)
+        logger.info("Created %s", model_path)
 
     def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_opset=None,
                  perf=None, fold_const=None):
         """Run complete test against backend."""
-        print(name)
+        logger.info("===================================")
+        logger.info("Running %s", name)
         self.perf = perf
 
         # get the model
@@ -216,7 +216,7 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
         else:
             model_path = self.local
             dir_name = os.path.dirname(self.local)
-        print("\tdownloaded", model_path)
+        logger.info("Downloaded to %s", model_path)
 
         input_names = list(self.input_names.keys())
         outputs = self.output_names
@@ -261,29 +261,24 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
 
             # run the model with tensorflow
             if self.skip_tensorflow:
-                print("\ttensorflow", "SKIPPED")
+                logger.info("TensorFlow SKIPPED")
             else:
                 tf_results = self.run_tensorflow(sess, inputs)
-                print("\ttensorflow", "OK")
+                logger.info("TensorFlow OK")
             model_proto = None
             try:
                 # convert model to onnx
                 onnx_graph = self.to_onnx(sess.graph, opset=opset, extra_opset=extra_opset,
                                           shape_override=shape_override, input_names=inputs.keys())
                 model_proto = onnx_graph.make_model("converted from tf2onnx")
-                new_model_proto = optimizer.optimize_graph(onnx_graph).make_model("optimized")
-                if new_model_proto:
-                    model_proto = new_model_proto
-                else:
-                    print("\tNON-CRITICAL, optimizers are not applied successfully")
-                print("\tto_onnx", "OK")
+                model_proto = optimizer.optimize_graph(onnx_graph).make_model("optimized")
+                logger.info("To_ONNX, OK")
                 if utils.is_debug_mode():
                     onnx_graph.dump_graph()
                 if onnx_file:
                     self.create_onnx_file(name, model_proto, inputs, onnx_file)
-            except Exception as ex:
-                tb = traceback.format_exc()
-                print("\tto_onnx", "FAIL", ex, tb)
+            except Exception:
+                logger.error("To_ONNX FAIL", exc_info=1)
 
         try:
             onnx_results = None
@@ -295,27 +290,25 @@ def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_ops
                 onnx_results = self.run_onnxruntime(name, model_proto, inputs)
             else:
                 raise ValueError("unknown backend")
-            print("\trun_onnx OK")
+            logger.info("Run_ONNX OK")
 
             try:
                 if self.skip_tensorflow:
-                    print("\tResults: skipped tensorflow")
+                    logger.info("Results: skipped tensorflow")
                 else:
                     if self.check_only_shape:
                         for tf_res, onnx_res in zip(tf_results, onnx_results):
                             np.testing.assert_array_equal(tf_res.shape, onnx_res.shape)
                     else:
                         for tf_res, onnx_res in zip(tf_results, onnx_results):
                             np.testing.assert_allclose(tf_res, onnx_res, rtol=self.rtol, atol=self.atol)
-                    print("\tResults: OK")
+                    logger.info("Results: OK")
                 return True
-            except Exception as ex:
-                tb = traceback.format_exc()
-                print("\tResults", ex, tb)
+            except Exception:
+                logger.error("Results", exc_info=1)
 
-        except Exception as ex:
-            tb = traceback.format_exc()
-            print("\trun_onnx", "FAIL", ex, tb)
+        except Exception:
+            logger.error("Run_ONNX FAIL", exc_info=1)
 
         return False
 
@@ -379,7 +372,7 @@ def main():
     Test.target = args.target
     tests = tests_from_yaml(args.config)
     if args.list:
-        print(sorted(tests.keys()))
+        logger.info(sorted(tests.keys()))
         return 0
     if args.tests:
         test_keys = args.tests.split(",")
@@ -397,17 +390,17 @@ def main():
             ret = t.run_test(test, backend=args.backend, onnx_file=args.onnx_file,
                              opset=args.opset, extra_opset=args.extra_opset, perf=args.perf,
                              fold_const=args.fold_const)
-        except Exception as ex:
+        except Exception:
+            logger.error("Failed to run %s", test, exc_info=1)
             ret = None
-            tb = traceback.format_exc()
-            print(ex, tb)
         finally:
             if not utils.is_debug_mode():
                 utils.delete_directory(TEMP_DIR)
         if not ret:
             failed += 1
 
-    print("=== RESULT: {} failed of {}, backend={}".format(failed, count, args.backend))
+    logger.info("===================================")
+    logger.info("RESULT: %s failed of %s, backend=%s", failed, count, args.backend)
 
     if args.perf:
         with open(args.perf, "w") as f:

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,16 +128,17 @@ 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("")
+    model_proto = GraphUtil.optimize_model_proto(model_proto)
 
     # write onnx graph
+    logger.info("")
+    logger.info("Successfully converted TensorFlow model %s to ONNX", model_path)
     if args.output:
         utils.save_protobuf(args.output, model_proto)
-        print("\nComplete successfully, the onnx model is generated at " + args.output)
+        logger.info("ONNX model is saved at %s", args.output)
+    else:
+        logger.info("To export ONNX model to file, please run with `--output` option")
 
 
 if __name__ == "__main__":

tf2onnx/graph.py

@@ -12,8 +12,6 @@
 import collections
 import copy
 import logging
-import sys
-import traceback
 import six
 import numpy as np
 
@@ -23,7 +21,6 @@
 from tf2onnx import optimizer
 from tf2onnx.schemas import get_schema
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -1039,18 +1036,19 @@ def extract_sub_graph_nodes(self, outputs_name, input_checker=None, ignore_unuse
 
     def delete_unused_nodes(self, outputs_name):
         """Delete nodes not in subgraph ending with output_names."""
-        if outputs_name:
-            # we need keep those placeholders that are used as input of Loop's body graph.
-            # some of them are not used in the graph, but still need be there to keep the graph complete.
-            related_nodes = self.extract_sub_graph_nodes(outputs_name, ignore_unused_placeholder=False)
-            for node in related_nodes:
-                attr_body_graphs = node.get_body_graphs()
-                if attr_body_graphs:
-                    for _, body_graph in attr_body_graphs.items():
-                        body_graph.delete_unused_nodes(body_graph.outputs)
-            self.reset_nodes(related_nodes)
-        else:
-            print("WARNING: outputs not specified, delete_unused_nodes not taking effect.")
+        if not outputs_name:
+            logger.debug("Outputs not specified, delete_unused_nodes not taking effect.")
+            return
+
+        # we need keep those placeholders that are used as input of Loop's body graph.
+        # some of them are not used in the graph, but still need be there to keep the graph complete.
+        related_nodes = self.extract_sub_graph_nodes(outputs_name, ignore_unused_placeholder=False)
+        for node in related_nodes:
+            attr_body_graphs = node.get_body_graphs()
+            if attr_body_graphs:
+                for _, body_graph in attr_body_graphs.items():
+                    body_graph.delete_unused_nodes(body_graph.outputs)
+        self.reset_nodes(related_nodes)
 
 
 class GraphUtil(object):
@@ -1066,7 +1064,7 @@ def optimize_model_proto(onnx_model_proto):
 
         Returns:
             model proto after optimization, if optimizer run successfully
-            or None, if exceptions happens
+            or onnx_model_proto, if exceptions happens
         """
         try:
             kwargs = GraphUtil.get_onnx_model_properties(onnx_model_proto)
@@ -1080,12 +1078,10 @@ def optimize_model_proto(onnx_model_proto):
                 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
+            # sometimes, onnx shape inference will fail for some reason,
+            # return onnx_model_proto for this case
+            logger.warning("Failed to optimize model proto", exc_info=1)
+            return onnx_model_proto
 
     @staticmethod
     def get_onnx_model_properties(onnx_model_proto):

tf2onnx/logging.py

@@ -47,15 +47,15 @@ def basicConfig(**kwargs):  # pylint: disable=invalid-name, function-redefined
     set_tf_verbosity(_logging.getLogger().getEffectiveLevel())
 
 
-_VERBOSITY_TO_LEVEL = [INFO, VERBOSE, DEBUG]
+_LOG_LEVELS = [FATAL, ERROR, WARNING, INFO, VERBOSE, DEBUG]
 
 
-def get_verbosity_level(verbosity, default_level=INFO):
+def get_verbosity_level(verbosity, base_level=INFO):
     """ If verbosity is specified, return corresponding level, otherwise, return default_level. """
     if verbosity is None:
-        return default_level
-    verbosity = min(max(0, verbosity), len(_VERBOSITY_TO_LEVEL) - 1)
-    return _VERBOSITY_TO_LEVEL[verbosity]
+        return base_level
+    verbosity = min(max(0, verbosity) + _LOG_LEVELS.index(base_level), len(_LOG_LEVELS) - 1)
+    return _LOG_LEVELS[verbosity]
 
 
 def set_level(level):
@@ -66,6 +66,9 @@ def set_level(level):
 
 def set_tf_verbosity(level):
     """ Set TF logging verbosity."""
+    # TF log is too verbose, adjust it
+    level = WARNING if level >= VERBOSE else level
+
     tf.logging.set_verbosity(level)
 
     # TF_CPP_MIN_LOG_LEVEL:

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),
+    ("optimize_transpose", TransposeOptimizer),
+    ("fold_constants", ConstFoldOptimizer),
+    # merge_duplication should be used after optimize_transpose
+    # for optimize_transpose may have some trans nodes that can be merge
+    ("merge_duplication", MergeDuplicatedNodesOptimizer),
+    ("remove_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