Merge pull request #329 from lucienwang1009/auto_infer_shape

auto infer shape and dtype for onnx nodes

Author: nbcsm

tests/backend_test_base.py

@@ -69,7 +69,7 @@ def run_onnxruntime(self, model_path, inputs, output_names):
         results = m.run(output_names, inputs)
         return results
 
-    def _run_backend(self, g, outputs, input_dict):
+    def run_backend(self, g, outputs, input_dict):
         model_proto = g.make_model("test")
         model_path = self.save_onnx_model(model_proto, input_dict)
 
@@ -133,7 +133,7 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
             g = process_tf_graph(sess.graph, opset=self.config.opset, output_names=output_names_with_port,
                                  target=self.config.target, **process_args)
             g = optimizer.optimize_graph(g)
-            actual = self._run_backend(g, output_names_with_port, onnx_feed_dict)
+            actual = self.run_backend(g, output_names_with_port, onnx_feed_dict)
 
         for expected_val, actual_val in zip(expected, actual):
             if check_value:

tests/test_onnx_shape_inference.py

@@ -15,11 +15,11 @@
 import six
 import numpy as np
 
-from onnx import helper, numpy_helper, shape_inference, OperatorSetIdProto, AttributeProto
+from onnx import helper, numpy_helper, shape_inference, OperatorSetIdProto, AttributeProto, TensorProto
 from tf2onnx import utils, __version__
 from tf2onnx.utils import port_name, find_opset
 from tf2onnx import optimizer
-from tf2onnx.schemas import get_schema
+from tf2onnx.schemas import get_schema, infer_onnx_shape_dtype
 
 logger = logging.getLogger(__name__)
 
@@ -419,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):
+                  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 = {}
@@ -474,6 +474,9 @@ 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 infer_shape_dtype:
+            self.update_node_shape_dtype(node, override=False)
+
         self._nodes.append(node)
         return node
 
@@ -533,6 +536,66 @@ def reset_nodes(self, ops):
         self._dtypes = remained_dtypes
         self._output_shapes = remained_shapes
 
+    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 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.
+        initializers = []
+        for i, inp in enumerate(node.inputs):
+            if not inp:
+                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]
+                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]
+
+        shapes, dtypes = infer_onnx_shape_dtype(node, self._opset, input_shapes, input_dtypes, initializers)
+        if not shapes or not dtypes:
+            return
+
+        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:
+                existing_dtype = self.get_dtype(output)
+                if existing_dtype is not None and existing_dtype != dtype:
+                    if override:
+                        logger.warning("Override dtype of %s from %s to %s", output, existing_dtype, dtype)
+                    else:
+                        dtype = existing_dtype
+                self.set_dtype(output, dtype)
+                logger.debug("Set dtype of [%s] to %s", output, dtype)
+
+            if shape is None:
+                logger.debug("Inferred shape for [%s, type: %s] is None, SKIP", node.name, node.type)
+            else:
+                existing_shape = self.get_shape(output)
+                if existing_shape is not None and not utils.are_shapes_equal(existing_shape, shape):
+                    if override:
+                        logger.warning("Override shape of %s from %s to %s", output, existing_shape, shape)
+                    else:
+                        shape = existing_shape
+                self.set_shape(output, shape)
+                logger.debug("Set shape of [%s] to %s", output, shape)
+
     def update_proto(self):
         """Update the onnx protobuf from out internal Node structure."""
         for node in self._nodes:

tf2onnx/graph.py

@@ -34,6 +34,7 @@ def get_inputs_for_current_iteration(g, input_id, iter_index):
 def create_loop_body_graph(parent_g, gather_input_ids, output_data_type, output_shape, trip_count_input_ids,
                            rank, loop_name):
     g = parent_g.create_new_graph_with_same_config()
+    g.parent_graph = parent_g
     iter_name = utils.make_name("i")
     cond_name = utils.make_name("cond")
     fake_var_name = utils.make_name("fake_var")
@@ -117,6 +118,7 @@ def create_if_op(g, input_ids, output_data_type, output_shape):
 
 def create_body_graph_for_if_branch(parent_g, data_type, output_shape, chosen_cur_cond_val_out_name, op_name):
     g = parent_g.create_new_graph_with_same_config()
+    g.parent_graph = parent_g
     name = utils.make_name("Identity")
     g.make_node(
         'Identity',
@@ -204,14 +206,15 @@ def make_range_non_const(ctx, start, limit, delta, output, scope_name, shape, dt
 
     # body
     g = ctx.create_new_graph_with_same_config()
-    g.make_node("Identity", ["cond"], outputs=["cond_out"])
-    g.make_node("Add", ["prev", delta], outputs=["current"], name=utils.make_name("add"))
-    g.make_node("Identity", ["prev"], outputs=["range"])
-
+    g.parent_graph = ctx
     g.add_graph_input("i", TensorProto.INT64, [])
     g.add_graph_input("cond", TensorProto.BOOL, [])
     g.add_graph_input("prev", dtype, [])
 
+    g.make_node("Identity", ["cond"], outputs=["cond_out"])
+    g.make_node("Add", ["prev", delta], outputs=["current"], name=utils.make_name("add"))
+    g.make_node("Identity", ["prev"], outputs=["range"])
+
     g.add_graph_output("cond_out", TensorProto.BOOL, [])
     g.add_graph_output("current", dtype, [])
     g.add_graph_output("range", dtype, [])
@@ -274,8 +277,9 @@ def version_8(cls, ctx, node, **kwargs):
         input_shape = ctx.get_shape(node.input[0])
 
         g = ctx.create_new_graph_with_same_config()
-        g.make_node('Identity', ['X'], outputs=['Y'])
+        g.parent_graph = ctx
         g.add_graph_input('X', input_dtype, input_shape[2:])
+        g.make_node('Identity', ['X'], outputs=['Y'])
         g.add_graph_output('Y', input_dtype, input_shape[2:])
 
         node.set_body_graph_as_attr("body", g)

tf2onnx/onnx_opset/controlflow.py

@@ -110,8 +110,8 @@ def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
                 shape_name = utils.make_name(node.name)
                 ctx.make_const(shape_name, np.array(new_kernel_shape, dtype=np.int64))
                 input_name = node.input[1]
-                reshape = ctx.insert_new_node_on_input(node, "Reshape", input_name)
-                reshape.input.append(shape_name)
+                reshape = ctx.make_node("Reshape", [input_name, shape_name])
+                ctx.replace_input(node, input_name, reshape.output[0])
                 reshape.skip_conversion = True
             ctx.set_shape(reshape.output[0], new_kernel_shape)
 
@@ -348,8 +348,8 @@ def version_7(cls, ctx, node, **kwargs):
                 shape_name = utils.make_name(node.name)
                 ctx.make_const(shape_name, np.array(new_broadcast_shape, dtype=np.int64))
                 op_name = node.input[1]
-                reshape_node = ctx.insert_new_node_on_input(node, "Reshape", op_name)
-                reshape_node.input.append(shape_name)
+                reshape_node = ctx.make_node("Reshape", [op_name, shape_name])
+                ctx.replace_input(node, op_name, reshape_node.output[0])
                 ctx.set_shape(reshape_node.output[0], new_broadcast_shape)
 
 

tf2onnx/onnx_opset/nn.py

@@ -283,15 +283,16 @@ def _make_gathernd_inner_loop(ctx, params, index, dtype):
 
     # body graph creation
     g = ctx.create_new_graph_with_same_config()
+    g.add_graph_input(trip_name, TensorProto.INT64, [])
+    g.add_graph_input(cond_name, TensorProto.BOOL, [])
+    g.add_graph_input(cur_name, dtype, [])
+    g.parent_graph = ctx
+
     index_i = g.make_node("Gather", [index.output[0], trip_name], attr={"axis": 0})
     gather = g.make_node("Gather", [cur_name, index_i.output[0]], attr={"axis": 0})
     g.make_node("Squeeze", [gather.output[0]], attr={"axes": [0]}, outputs=[result_name])
     g.make_node("Identity", [cond_name], outputs=[cond_out_name])
 
-    g.add_graph_input(trip_name, TensorProto.INT64, [])
-    g.add_graph_input(cond_name, TensorProto.BOOL, [])
-    g.add_graph_input(cur_name, dtype, [])
-
     g.add_graph_output(cond_out_name, TensorProto.BOOL, [])
     g.add_graph_output(result_name, dtype, [])
 
@@ -337,6 +338,11 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params, shapes, dt
     dummy_out_name = utils.make_name("dummy_out")
     result_name = utils.make_name("res")
 
+    g.add_graph_input(trip_name, TensorProto.INT64, [])
+    g.add_graph_input(cond_name, TensorProto.BOOL, [])
+    g.add_graph_input(dummy_name, t_params, [])
+    g.parent_graph = ctx
+
     index = g.make_node("Gather", [flatten_indices.output[0], trip_name], attr={"axis": 0})
     index_squeeze = g.make_node("Squeeze", [index.output[0]], attr={"axes": [0]})
     # inner loop to gather result
@@ -345,10 +351,6 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params, shapes, dt
     g.make_node("Identity", [dummy_name], outputs=[dummy_out_name])
     g.make_node("Identity", [inner_loop.output[0]], outputs=[result_name])
 
-    g.add_graph_input(trip_name, TensorProto.INT64, [])
-    g.add_graph_input(cond_name, TensorProto.BOOL, [])
-    g.add_graph_input(dummy_name, t_params, [])
-
     g.add_graph_output(cond_out_name, TensorProto.BOOL, [])
     g.add_graph_output(dummy_out_name, t_params, [])
     g.add_graph_output(result_name, t_params, [])

tf2onnx/onnx_opset/tensor.py

@@ -9,10 +9,15 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import logging
+import copy
 from collections import defaultdict, OrderedDict
-from onnx import defs
+from onnx import defs, helper, TensorProto, OperatorSetIdProto, shape_inference
 
 from . import constants
+from . import utils
+
+logger = logging.getLogger(__name__)
 
 
 class OnnxOpSchema(object):
@@ -112,3 +117,75 @@ def get_max_supported_opset_version(domain=None):
     """Get max supported opset version by current onnx package given a domain."""
     domain = domain or constants.ONNX_DOMAIN
     return _domain_opset_versions.get(domain, 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.
+    """
+
+    def build_onnx_op(node):
+        """Build onnx op"""
+        onnx_node = helper.make_node(node.type, node.input, node.output, name=node.name)
+        # deal with attributes
+        attr = []
+        attr_graphs = node.get_body_graphs()
+        if attr_graphs:
+            for attr_name, sub_graph in attr_graphs.items():
+                copied_sub_graph = copy.deepcopy(sub_graph)
+                graph_proto = copied_sub_graph.make_graph("graph for " + node.name + " " + attr_name)
+                attr.append(helper.make_attribute(attr_name, graph_proto))
+        attr.extend([a for a in node.attr_onnx.values()])
+        if attr:
+            onnx_node.attribute.extend(attr)
+        return onnx_node
+
+    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_proto = helper.make_graph([build_onnx_op(node)], "infer-graph", inputs, outputs, initializer=initializers)
+    imp = OperatorSetIdProto()
+    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
+
+    shapes = {}
+    dtypes = {}
+    for output in inferred_model.graph.output:
+        tensor_type = output.type.tensor_type
+        if tensor_type.HasField("elem_type"):
+            dtypes[output.name] = tensor_type.elem_type
+        else:
+            dtypes[output.name] = TensorProto.UNDEFINED
+        # 0 in shapes of onnx means unknown which is -1 in our convertor
+        if tensor_type.HasField("shape"):
+            shapes[output.name] = [
+                dim.dim_value if dim.dim_value != 0 else utils.ONNX_UNKNOWN_DIMENSION for dim in tensor_type.shape.dim
+            ]
+        else:
+            shapes[output.name] = None
+    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/schemas.py

@@ -244,7 +244,14 @@ def make_onnx_inputs_outputs(name, elem_type, shape, **kwargs):
        elem_type,  # type: TensorProto.DataType
        shape,  # type: Optional[Sequence[int]]
     """
-    return helper.make_tensor_value_info(name, elem_type, make_onnx_shape(shape), **kwargs)
+    if elem_type is None:
+        elem_type = onnx_pb.TensorProto.UNDEFINED
+    return helper.make_tensor_value_info(
+        name,
+        elem_type,
+        make_onnx_shape(shape),
+        **kwargs
+    )
 
 
 def find_opset(opset):
@@ -308,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)
+                               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():
@@ -327,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])
+                           shapes=[shape], dtypes=[dtype], infer_shape_dtype=False)
         new_output_names.append(node.output[0])
     g.outputs = new_output_names
     return g