Merge pull request #487 from zhijxu-MS/push_branch

enhance shape inference, fix bugs

Author: nbcsm

tests/common.py

@@ -11,6 +11,7 @@
 
 from distutils.version import LooseVersion
 from parameterized import parameterized
+import numpy as np
 from tf2onnx import constants, logging, utils
 
 __all__ = [
@@ -280,7 +281,8 @@ def check_onnxruntime_incompatibility(op):
 def validate_const_node(node, expected_val):
     if node.is_const():
         node_val = node.get_tensor_value()
-        return node_val == expected_val
+        np.testing.assert_allclose(expected_val, node_val)
+        return True
     return False
 
 

tf2onnx/graph_builder.py

@@ -70,6 +70,11 @@ def make_slice(self, kwargs, name=None, shapes=None, dtypes=None):
         while inputs[-1] == "":
             inputs = inputs[:-1]
 
+        if self.graph.opset >= 10:
+            dtype = self.graph.get_dtype(inputs[1])
+            for input_data in inputs[1:]:
+                utils.make_sure(dtype == self.graph.get_dtype(input_data), "dtype should be same")
+
         return self.graph.make_node(op_type="Slice", inputs=inputs, attr=attr, name=name,
                                     outputs=outputs, shapes=shapes, dtypes=dtypes).output[0]
 

tf2onnx/onnx_opset/tensor.py

@@ -534,11 +534,12 @@ def version_4(cls, ctx, node, **kwargs):
             attr = node.get_attr(attr_name)
             if attr is not None and attr.i != 0:
                 raise ValueError("StridedSlice: attribute " + attr_name + " not supported")
-        input_shape = ctx.get_shape(node.input[0])
-        begin = node.inputs[1].get_tensor_value(as_list=False)
-        end = node.inputs[2].get_tensor_value(as_list=False)
-        strides = node.inputs[3].get_tensor_value(as_list=False)
-        max_size = np.iinfo(begin.dtype).max
+        onnx_dtype = ctx.get_dtype(node.input[1])
+        np_dtype = utils.ONNX_TO_NUMPY_DTYPE[onnx_dtype]
+        max_size = np.iinfo(np_dtype).max
+        begin = node.inputs[1].get_tensor_value()
+        end = node.inputs[2].get_tensor_value()
+        strides = node.inputs[3].get_tensor_value()
         end_mask = node.get_attr("end_mask")
         end_mask = end_mask.i if end_mask is not None else 0
         begin_mask = node.get_attr("begin_mask")

tf2onnx/rewriter/bigru_rewriter.py

@@ -13,7 +13,7 @@
 import logging
 import numpy as np
 from tf2onnx import utils
-from tf2onnx.rewriter.rnn_utils import is_reverse_op
+from tf2onnx.utils import is_reverse_op
 from tf2onnx.rewriter.bilstm_rewriter import slice_bilstm_for_original_lstm_consumers,\
      get_reverse_nodes_after_y_output, get_np_val_for_const, _process_single_init_node
 

tf2onnx/rewriter/bilstm_rewriter.py

@@ -13,7 +13,7 @@
 import logging
 import numpy as np
 from tf2onnx import utils
-from tf2onnx.rewriter.rnn_utils import is_reverse_op
+from tf2onnx.utils import is_reverse_op
 from tf2onnx.graph_builder import GraphBuilder
 
 logger = logging.getLogger(__name__)

tf2onnx/rewriter/loop_rewriter_base.py

@@ -12,8 +12,8 @@
 from collections import OrderedDict
 from tf2onnx import utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
-from tf2onnx.rewriter.rnn_utils import is_loopcond_op, is_tensor_array_op
-from tf2onnx.rewriter.rnn_utils import is_tensor_array_gather_op, is_tensor_array_write_op
+from tf2onnx.utils import is_loopcond_op, is_tensor_array_op
+from tf2onnx.utils import is_tensor_array_gather_op, is_tensor_array_write_op
 from tf2onnx.rewriter.rnn_utils import REWRITER_RESULT
 from tf2onnx.utils import TensorValueInfo
 

tf2onnx/rewriter/lstm_rewriter.py

@@ -13,8 +13,8 @@
 import numpy as np
 from tf2onnx import utils
 from tf2onnx.graph_builder import GraphBuilder
-from tf2onnx.rewriter.rnn_utils import RNNUnitType, RnnWeight, \
-    is_concat_op, is_slice_op, get_weights_from_const_node
+from tf2onnx.rewriter.rnn_utils import RNNUnitType, RnnWeight, get_weights_from_const_node
+from tf2onnx.utils import is_concat_op, is_slice_op
 
 from tf2onnx.rewriter.unit_rnn_rewriter_base import UnitRnnRewriterBase
 

tf2onnx/rewriter/rnn_utils.py

@@ -263,35 +263,3 @@ def get_weights_from_const_node(g, node):
         return None
 
     return RnnWeight(node, val, dtype)
-
-
-def is_reverse_op(op):
-    return op.type in ("ReverseV2", "ReverseSequence")
-
-
-def is_concat_op(op):
-    return op.type in ("Concat", "ConcatV2", "ConcatV3")
-
-
-def is_tensor_array_gather_op(op):
-    return op.type in ("TensorArrayGatherV2", "TensorArrayGatherV3")
-
-
-def is_tensor_array_write_op(op):
-    return op.type in ("TensorArrayWriteV2", "TensorArrayWriteV3")
-
-
-def is_tensor_array_op(op):
-    return op.type in ("TensorArrayV2", "TensorArrayV3")
-
-
-def is_loopcond_op(op):
-    return op.type == "LoopCond"
-
-
-def is_select_op(op):
-    return op.type == "Select"
-
-
-def is_slice_op(op):
-    return op.type == "Slice"

tf2onnx/rewriter/unit_rnn_rewriter_base.py

@@ -13,12 +13,11 @@
 from tf2onnx.graph_builder import GraphBuilder
 from tf2onnx.rewriter.loop_rewriter_base import LoopRewriterBase, Context
 from tf2onnx.rewriter.rnn_utils import REWRITER_RESULT, get_pattern, \
-    get_rnn_scope_name, parse_rnn_loop, is_select_op, is_tensor_array_write_op, \
-    seq_len_pattern
+    get_rnn_scope_name, parse_rnn_loop, seq_len_pattern
+from tf2onnx.utils import is_select_op, is_tensor_array_write_op
 from tf2onnx.graph_matcher import GraphMatcher
 
 
-
 logger = logging.getLogger(__name__)
 
 

tf2onnx/shape_inference.py

@@ -9,13 +9,13 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 import logging
+import numpy as np
 from onnx import onnx_pb
 from tf2onnx import utils
 
 # pylint: disable=logging-not-lazy,missing-docstring,consider-swap-variables
 
 
-
 logger = logging.getLogger(__name__)
 
 direct_ops = [
@@ -115,43 +115,19 @@ def infer_shape_for_node(g, node):
             return False
         return set_shape_from_input(g, shape_node.input[0], node.output[0])
 
-    if node.type == "ConcatV2":
-        axis_node = node.inputs[-1]
-        if not axis_node.is_const():
-            return False
-
-        axis = axis_node.get_tensor_value()
-        val = 0
-        data_inputs = node.input[:-1]
-        for i in data_inputs:
-            s = g.get_shape(i)
-            if s is None:
-                return False
-
-            if s[axis] == -1:
-                val = -1
-                break
-            val += s[axis]
-
-        s1 = g.get_shape(node.input[0])
-        if axis < 0:
-            axis += len(s1)
-        new_shape = s1[:axis] + [val]
-        if axis < len(s1) - 1:
-            new_shape += s1[axis + 1:]
-
-        g.set_shape(node.output[0], new_shape)
-        logger.debug("set ConcatV2 node [%s] with new shape %s", node.output[0], new_shape)
-        return True
-
     if node.type == "Gather":
         # uses the follwing link to know how to infer shape of output
         # https://www.tensorflow.org/api_docs/python/tf/gather
         shape_params = g.get_shape(node.input[0])
         shape_indices = g.get_shape(node.input[1])
-        axis = node.input[2].get_tensor_value()
+        # gather can only have 2 inputs
+        # https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/gather.html
+        if len(node.input) == 3:
+            axis = node.input[2].get_tensor_value()
+        else:
+            axis = 0
 
-        shape = shape_params[:axis] + shape_indices + shape_indices[axis + 1:]
+        shape = shape_params[:axis] + shape_indices + shape_params[axis + 1:]
         g.set_shape(node.output[0], shape)
         return True
 
@@ -194,6 +170,29 @@ def infer_shape_for_node(g, node):
         logger.debug("set [%s] with new shape %s", node.output[0], new_shape)
         return True
 
+    if node.type == "Unpack":
+        input_shape = g.get_shape(node.input[0])
+        if input_shape is None:
+            return False
+
+        axis = node.get_attr("axis").i
+        axis = axis if axis >= 0 else axis + len(input_shape)
+        # the link below says that the rank of output is "rank(input) -1",
+        # from this statement "num" must equal to input_shape[axis], and if not tf will throw a runtime error
+        # https://www.tensorflow.org/api_docs/python/tf/unstack
+        new_shape = input_shape[:axis] + input_shape[axis + 1:]
+        for output in node.output:
+            g.set_shape(output, new_shape)
+            logger.debug("set %s node [%s] with new shape %s", node.type, output, new_shape)
+        return True
+
+    if node.type in ["Minimum", "Maximum"]:
+        # ops that are elementwise and support broadcasting
+        input_shapes = [g.get_shape(node) for node in node.input]
+        new_shape = broadcast_shape_inference(*input_shapes)
+        g.set_shape(node.output[0], new_shape)
+        return True
+
     return False
 
 
@@ -213,6 +212,36 @@ def infer_input_shapes(g, node):
 
 
 def infer_output_shapes_with_partial_inputs(g, node):
+    # output shape of concat op: only the dim val of concatenated dim will be changed
+    # so only partial(at least one) input shapes need to be known to infer output shape of concat node
+    if utils.is_concat_op(node):
+        data_inputs = node.input[:-1]
+        input_shapes = [g.get_shape(node) for node in data_inputs]
+        input_shapes = [shape for shape in input_shapes if shape is not None]
+        if not input_shapes:
+            logger.debug("all input shapes of concat node %s are None, can't infer its output shape", node.name)
+            return False
+
+        new_shape = input_shapes[0]
+        axis_node = node.inputs[-1]
+        rank = len(new_shape)
+        if not axis_node.is_const():
+            g.set_shape(node.output[0], [-1] * rank)
+            return True
+
+        axis = axis_node.get_tensor_value()
+        axis = axis if axis >= 0 else axis + rank
+        new_shape[axis] = -1
+        if len(input_shapes) == len(data_inputs):  # all input shapes are known
+            concat_dim_vals = list(np.array(input_shapes)[:, axis])
+            # only when inputs' shape are known, then val of concat dim can be calculated
+            if concat_dim_vals.count(-1) == 0:
+                new_shape[axis] = sum(concat_dim_vals)
+
+        g.set_shape(node.output[0], new_shape)
+        logger.debug("set Concat node [%s] with new shape %s", node.output[0], new_shape)
+        return True
+
     if node.type == "Merge":
         s1 = g.get_shape(node.input[0])
         s2 = g.get_shape(node.input[1])