fix review comments

Author: nbcsm

tf2onnx/graph.py

@@ -1063,18 +1063,25 @@ def optimize_model_proto(onnx_model_proto):
         """Optimize the model proto, for example: eliminating all useless Transpose pairs.
 
         Returns:
-            model proto after optimization
+            model proto after optimization, if optimizer run successfully
+            or onnx_model_proto, if exceptions happens
         """
-        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
+        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,
+            # 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/optimizer/__init__.py

@@ -17,12 +17,12 @@
 
 # optimizer sequence need to be considered carefully
 _optimizers = OrderedDict([
-    ("reduce_transpose", TransposeOptimizer),
+    ("optimize_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 should be used after optimize_transpose
+    # for optimize_transpose may have some trans nodes that can be merge
     ("merge_duplication", MergeDuplicatedNodesOptimizer),
-    ("reduce_identity", IdentityOptimizer),
+    ("remove_identity", IdentityOptimizer),
 ])
 
 

tf2onnx/optimizer/optimizer_base.py

@@ -26,7 +26,7 @@ def is_debug_mode(self):
         return utils.is_debug_mode()
 
     def optimize(self, graph):
-        """ Optimize graph, return optimized graph """
+        """ Optimize graph, return optimized graph. """
         before = graph.dump_node_statistics()
 
         graph = self._optimize(graph)
@@ -38,6 +38,7 @@ def optimize(self, graph):
         return graph
 
     def _optimize(self, graph):
+        """ Derived class should override this function. """
         raise NotImplementedError
 
     @staticmethod

tf2onnx/optimizer/transpose_optimizer.py

@@ -129,10 +129,9 @@ def merge_duplicated_transposes(self):
         # dangling transpose nodes can be deleted
         graph.delete_unused_nodes(graph.outputs)
 
-    def optimize(self, graph):
+    def _optimize(self, graph):
         self._g = graph
         self.pre_optimize_action()
-        previous_counter = self._g.dump_node_statistics()
         no_action = False
         iteration_cnt = 0
         while not no_action:
@@ -161,10 +160,6 @@ def optimize(self, graph):
 
         self.merge_duplicated_transposes()
         self.post_optimize_action()
-
-        current_counter = self._g.dump_node_statistics()
-        transpose_cnt = current_counter["Transpose"]
-        self._print_stat_diff(previous_counter, current_counter)
         return self._g
 
     def _initialize_handlers(self):