resolve comments

Author: lucienwang1009

tests/test_shape_inference.py renamed to tests/test_onnx_shape_inference.py

@@ -31,25 +31,12 @@ class ShapeInferenceTests(Tf2OnnxBackendTestBase):
     def _run_test_case(self, graph, feed_dict):
         """Run model with onnxruntime and compare results' shape with internal shape inference."""
         outputs = graph.outputs
-        results = []
-        raised = False
-        try:
-            results = self.run_backend(graph, outputs, feed_dict)
-        except Exception as ex:  # pylint: disable=broad-except
-            self.logger.error(ex)
-            raised = True
-        self.assertFalse(raised)
-
-        raised = False
+        results = self.run_backend(graph, outputs, feed_dict)
+
         for actual, inferred in zip(results, outputs):
             actual_shape = actual.shape
             inferred_shape = tuple(graph.get_shape(inferred))
-            try:
-                utils.merge_shapes(actual_shape, inferred_shape)
-            except Exception as ex:  # pylint: disable=broad-except
-                self.logger.error(ex)
-                raised = True
-            self.assertFalse(raised)
+            self.assertTrue(utils.are_shapes_compatible(actual_shape, inferred_shape))
 
             actual_dtype = actual.dtype
             inferred_dtype = utils.ONNX_TO_NUMPY_DTYPE[graph.get_dtype(inferred)]
@@ -418,6 +405,9 @@ def test_override_shape(self):
         output_name = utils.make_name("output")
         graph._output_shapes[output_name] = [-1, -1, 2, 3]  # pylint: disable=protected-access
         node = graph.make_node("Transpose", [INPUT1], attr={"perm": [1, 0, 2, 3]}, outputs=[output_name])
+
+        graph.update_node_shape_dtype(node, override=True)
+
         graph.add_graph_output(node.output[0])
         self._run_test_case(graph, self._generate_random_inputs(inputs, shapes, dtypes))
 

tf2onnx/graph.py

@@ -12,12 +12,10 @@
 import collections
 import copy
 import logging
-import traceback
 import six
 import numpy as np
 
 from onnx import helper, numpy_helper, shape_inference, OperatorSetIdProto, AttributeProto, TensorProto
-from tf2onnx import constants
 from tf2onnx import utils, __version__
 from tf2onnx.utils import port_name, find_opset
 from tf2onnx import optimizer
@@ -421,7 +419,7 @@ def make_const(self, name, np_val, skip_conversion=False, raw=True):
         return node
 
     def make_node(self, op_type, inputs, attr=None, output_count=1, outputs=None, skip_conversion=True,
-                  op_name_scope=None, name=None, shapes=None, dtypes=None, domain=None, auto_infer_shape_dtype=True):
+                  op_name_scope=None, name=None, shapes=None, dtypes=None, domain=None, infer_shape_dtype=True):
         """Make a new onnx node in the graph"""
         if attr is None:
             attr = {}
@@ -476,8 +474,8 @@ def make_node(self, op_type, inputs, attr=None, output_count=1, outputs=None, sk
             for i in range(output_count):
                 self.set_dtype(node.output[i], dtypes[i])
 
-        if (not shapes or not dtypes) and auto_infer_shape_dtype:
-            self.update_node_shape_dtype(node, override=True)
+        if (not shapes or not dtypes) and infer_shape_dtype:
+            self.update_node_shape_dtype(node, override=False)
 
         self._nodes.append(node)
         return node
@@ -538,44 +536,42 @@ def reset_nodes(self, ops):
         self._dtypes = remained_dtypes
         self._output_shapes = remained_shapes
 
-    def update_node_shape_dtype(self, node, override=True):
-        """try the best to infer shapes and dtypes for outputs of the node"""
+    def update_node_shape_dtype(self, node, override=False):
+        """Try the best to infer shapes and dtypes for outputs of the node,
+        by default, we respect TF shapes and dtypes.
+        """
         if node.is_const() or node.is_graph_input():
             return
         # NOTE: only support onnx node for now
-        if node.domain != constants.ONNX_DOMAIN:
+        if not utils.is_onnx_domain(node.domain):
             return
 
         logger.debug("Infer shape and dtype for [%s]", node.name)
         # NOTE: shape inference for some ops need the input values of the op, e.g., Reshape
         # op needs the "Shape" value to infer output shape.
-        initializer = []
+        initializers = []
         for i, inp in enumerate(node.inputs):
             if not inp:
-                logger.warning("[%s] infer a inexistent node: [%s], please check the code", node.name, node.input[i])
+                if logger.isEnabledFor(logging.VERBOSE):
+                    logger.warning(
+                        "[%s] infer a inexistent node: [%s], please check the code",
+                        node.name, node.input[i]
+                    )
                 continue
             if inp.is_const():
                 t = inp.get_attr("value")
                 tensor = helper.get_attribute_value(t)
                 tensor.name = inp.output[0]
-                initializer.append(tensor)
+                initializers.append(tensor)
 
         input_shapes = [self.get_shape(i) for i in node.input]
         input_dtypes = [self.get_dtype(i) for i in node.input]
 
-        dtypes = {}
-        shapes = {}
-        try:
-            shapes, dtypes = infer_onnx_shape_dtype(node, input_shapes, input_dtypes, self._opset, initializer)
-        except Exception:
-            tb = traceback.format_exc()
-            logger.warning("ONNX Failed to infer shapes and dtypes for [%s, type: %s]", node.name, node.type)
-            logger.warning("Inference error: %s", tb)
+        shapes, dtypes = infer_onnx_shape_dtype(node, self._opset, input_shapes, input_dtypes, initializers)
+        if not shapes or not dtypes:
             return
 
-        for output in node.output:
-            dtype = dtypes[output]
-            shape = shapes[output]
+        for output, shape, dtype in zip(node.output, shapes, dtypes):
             if dtype == TensorProto.UNDEFINED:
                 logger.debug("Inferred dtype for [%s, type: %s] is UNDEFINED, SKIP", node.name, node.type)
             else:

tf2onnx/schemas.py

@@ -9,13 +9,16 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import logging
 import copy
 from collections import defaultdict, OrderedDict
 from onnx import defs, helper, TensorProto, OperatorSetIdProto, shape_inference
 
 from . import constants
 from . import utils
 
+logger = logging.getLogger(__name__)
+
 
 class OnnxOpSchema(object):
     """Wrapper for Onnx schema."""
@@ -116,7 +119,7 @@ def get_max_supported_opset_version(domain=None):
     return _domain_opset_versions.get(domain, None)
 
 
-def infer_onnx_shape_dtype(node, input_shapes, input_dtypes, opset, initializer=None):
+def infer_onnx_shape_dtype(node, opset_version, input_shapes, input_dtypes, initializers=None):
     """
     Infer shapes and dtypes for outputs of the node.
     Sometimes, shape inference needs the values of node's inputs, so initializers are used.
@@ -138,20 +141,29 @@ def build_onnx_op(node):
             onnx_node.attribute.extend(attr)
         return onnx_node
 
-    shapes = {}
-    dtypes = {}
     inputs = []
     outputs = []
     for inp, shape, dtype in zip(node.input, input_shapes, input_dtypes):
         inputs.append(utils.make_onnx_inputs_outputs(inp, dtype, shape))
     for output in node.output:
         outputs.append(utils.make_onnx_inputs_outputs(output, TensorProto.UNDEFINED, None))
-    graph_def = helper.make_graph([build_onnx_op(node)], "infer-graph", inputs, outputs, initializer=initializer)
+    graph_proto = helper.make_graph([build_onnx_op(node)], "infer-graph", inputs, outputs, initializer=initializers)
     imp = OperatorSetIdProto()
-    imp.version = opset
-    model_def = helper.make_model(graph_def, opset_imports=[imp])
+    imp.version = opset_version
+    model_proto = helper.make_model(graph_proto, opset_imports=[imp])
+
+    inferred_model = None
+    try:
+        inferred_model = shape_inference.infer_shapes(model_proto)
+    except Exception:  # pylint: disable=broad-except
+        logger.warning(
+            "ONNX Failed to infer shapes and dtypes for [%s, type: %s]",
+            node.name, node.type, exc_info=1
+        )
+        return None, None
 
-    inferred_model = shape_inference.infer_shapes(model_def)
+    shapes = {}
+    dtypes = {}
     for output in inferred_model.graph.output:
         tensor_type = output.type.tensor_type
         if tensor_type.HasField("elem_type"):
@@ -165,4 +177,15 @@ def build_onnx_op(node):
             ]
         else:
             shapes[output.name] = None
-    return shapes, dtypes
+    output_shapes = []
+    output_dtypes = []
+    for output in node.output:
+        if output in shapes:
+            output_shapes.append(shapes[output])
+        else:
+            output_shapes.append(None)
+        if output in dtypes:
+            output_dtypes.append(dtypes[output])
+        else:
+            output_dtypes.append(TensorProto.UNDEFINED)
+    return output_shapes, output_dtypes

tf2onnx/utils.py

@@ -315,7 +315,7 @@ def construct_graph_from_nodes(parent_g, nodes, outputs, shapes, dtypes):
         all_outputs |= set(op.output)
 
         new_node = g.make_node(op.type, op.input, outputs=op.output, attr=op.attr, name=op.name,
-                               skip_conversion=op.skip_conversion, auto_infer_shape_dtype=False)
+                               skip_conversion=op.skip_conversion, infer_shape_dtype=False)
         body_graphs = op.graph.contained_graphs.pop(op.name, None)
         if body_graphs:
             for attr_name, body_graph in body_graphs.items():
@@ -334,7 +334,7 @@ def construct_graph_from_nodes(parent_g, nodes, outputs, shapes, dtypes):
     new_output_names = []
     for output, shape, dtype in zip(outputs, shapes, dtypes):
         node = g.make_node("Identity", inputs=[output], op_name_scope="sub_graph_ending_node",
-                           shapes=[shape], dtypes=[dtype], auto_infer_shape_dtype=False)
+                           shapes=[shape], dtypes=[dtype], infer_shape_dtype=False)
         new_output_names.append(node.output[0])
     g.outputs = new_output_names
     return g

tf2onnx/verbose_logging.py

@@ -32,7 +32,7 @@ def getLogger(name=None):  # pylint: disable=invalid-name, function-redefined
     return logger
 
 
-_SIMPLE_LOG_FORMAT = "%(levelname)s: %(message)s"
+_SIMPLE_LOG_FORMAT = "%(message)s"
 _VERBOSE_LOG_FORMAT = "%(asctime)s - %(levelname)s - %(name)s: %(message)s"