Merge pull request #350 from nbcsm/shape

shape inference support LogicalOr, relax shape check for cond rewriter

Author: nbcsm

tests/test_internals.py

@@ -16,8 +16,7 @@
 from onnx import helper
 
 import tensorflow as tf
-import tf2onnx
-import tf2onnx.utils
+from tf2onnx import utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.graph import GraphUtil
 from common import unittest_main
@@ -50,14 +49,13 @@ def onnx_pretty(g, args=None):
 
 
 class Tf2OnnxInternalTests(unittest.TestCase):
-
     def setUp(self):
         """Setup test."""
         # suppress log info of tensorflow so that result of test can be seen much easier
         os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
         tf.logging.set_verbosity(tf.logging.WARN)
 
-        tf2onnx.utils.INTERNAL_NAME = 1
+        utils.INTERNAL_NAME = 1
         arg = namedtuple("Arg", "input inputs outputs verbose")
         self._args0 = arg(input="test", inputs=[], outputs=["output:0"], verbose=False)
         self._args1 = arg(input="test", inputs=["input:0"], outputs=["output:0"], verbose=False)
@@ -142,8 +140,8 @@ def test_rewrite_subgraph(self):
         for match in match_results:
             input_node = match.get_op('input')
             output_node = match.get_op('output')
-            op_name = tf2onnx.utils.make_name("ReplacedOp")
-            out_name = tf2onnx.utils.port_name(op_name)
+            op_name = utils.make_name("ReplacedOp")
+            out_name = utils.port_name(op_name)
             new_node = g.make_node("Sub", inputs=input_node.input, outputs=[out_name], name=op_name)
             ops = g.replace_subgraph(ops, match, [], [output_node], [], [new_node])
         g.topological_sort(ops)
@@ -183,10 +181,26 @@ def test_cmdarg_parse(self):
         arg = "input/V-1_2:0,input/X:0[1,2,3],Y:1[4,5],Z:3,A:1,B"
         expected_inputs = ['input/V-1_2:0', 'input/X:0', 'Y:1', 'Z:3', 'A:1', 'B']
         expected_shape = {'Y:1': [4, 5], 'input/X:0': [1, 2, 3]}
-        inputs, shape_override = tf2onnx.utils.split_nodename_and_shape(arg)
+        inputs, shape_override = utils.split_nodename_and_shape(arg)
         self.assertEqual(expected_inputs, inputs)
         self.assertEqual(expected_shape, shape_override)
 
+    def test_shape_utils(self):
+        self.assertEqual(utils.merge_shapes(None, None), None)
+        self.assertEqual(utils.merge_shapes([], None), [])
+        self.assertEqual(utils.merge_shapes(None, [1, 2, 3]), [1, 2, 3])
+        self.assertEqual(utils.merge_shapes([1, 3], [None, 3]), [1, 3])
+        self.assertEqual(utils.merge_shapes([1, None, 3], (-1, 2, "unk")), [1, 2, 3])
+
+        self.assertTrue(utils.are_shapes_compatible(None, []))
+        self.assertTrue(utils.are_shapes_compatible([1, None, 3], (-1, 2, "unk")))
+        self.assertFalse(utils.are_shapes_compatible([1, 2, 3], (2, 3)))
+        self.assertFalse(utils.are_shapes_compatible([1, 2, 3], (4, 5, 6)))
+
+        self.assertTrue(utils.are_shapes_equal(None, None))
+        self.assertFalse(utils.are_shapes_equal(None, []))
+        self.assertTrue(utils.are_shapes_equal([1, 2, 3], (1, 2, 3)))
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/rewriter/cond_rewriter.py

@@ -16,6 +16,7 @@
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger("tf2onnx.rewriter.cond_rewriter_base")
 
+
 # pylint: disable=missing-docstring,unused-argument,broad-except
 
 class BranchType(Enum):
@@ -29,6 +30,7 @@ class BranchType(Enum):
 
 class CondBranchContext:
     """Context for each branch graph"""
+
     def __init__(self):
         self.output = []
         self.nodes = set()
@@ -37,12 +39,12 @@ def __init__(self):
 class CondContext:
     def __init__(self, cond_scope, pred_input, true_branch_context,
                  false_branch_context, switchs, merges):
-        self.cond_scope = cond_scope # name scope for this tf.cond
-        self.pred_input = pred_input # condition input
+        self.cond_scope = cond_scope  # name scope for this tf.cond
+        self.pred_input = pred_input  # condition input
         self.true_branch_context = true_branch_context
         self.false_branch_context = false_branch_context
         self.switchs = set(switchs)
-        self.merges = merges # list of merges in order
+        self.merges = merges  # list of merges in order
 
 
 class CondRewriter:
@@ -114,7 +116,7 @@ def _get_output_shape_dtype(self, cond_context):
             true_dtype = self.g.get_dtype(true_output)
             false_shape = self.g.get_shape(false_output)
             false_dtype = self.g.get_dtype(false_output)
-            if true_shape != false_shape:
+            if not utils.are_shapes_compatible(true_shape, false_shape):
                 raise RuntimeError(
                     "the shape of outputs {} and {} mismatch: {}, {}".format(
                         true_output,
@@ -132,7 +134,7 @@ def _get_output_shape_dtype(self, cond_context):
                         false_dtype
                     )
                 )
-            output_shapes.append(true_shape)
+            output_shapes.append(utils.merge_shapes(true_shape, false_shape))
             output_dtypes.append(true_dtype)
         return output_shapes, output_dtypes
 
@@ -243,11 +245,13 @@ def _trace_back_from_one_merge(self, merge_node):
         merge_input_1 = merge_node.input[0]
         merge_input_2 = merge_node.input[1]
         switchs = set()
+
         def stop_at_switch(node):
             if self._is_switch(node):
                 switchs.add(node)
                 return False
             return True
+
         branch_nodes_1 = self.g.extract_sub_graph_nodes(
             [merge_input_1],
             stop_at_switch

tf2onnx/shape_inference.py

@@ -36,11 +36,13 @@
     "GreaterEqual",
     "Less",
     "LogicalAnd",
+    "LogicalOr",
     "Mul",
     "RealDiv",
     "Sub"
 ]
 
+
 def infer_shape_for_graph(g):
     no_shape_updated = True
     while no_shape_updated:
@@ -135,7 +137,7 @@ def infer_shape_for_node(g, node):
             axis += len(s1)
         new_shape = s1[:axis] + [val]
         if axis < len(s1) - 1:
-            new_shape += s1[axis+1:]
+            new_shape += s1[axis + 1:]
 
         g.set_shape(node.output[0], new_shape)
         log.debug("set ConcatV2 node [%s] with new shape %s", node.output[0], new_shape)
@@ -148,7 +150,7 @@ def infer_shape_for_node(g, node):
         shape_indices = g.get_shape(node.input[1])
         axis = node.input[2].get_tensor_value()
 
-        shape = shape_params[:axis] + shape_indices + shape_indices[axis+1:]
+        shape = shape_params[:axis] + shape_indices + shape_indices[axis + 1:]
         g.set_shape(node.output[0], shape)
         return True
 

tf2onnx/utils.py

@@ -342,11 +342,6 @@ def tf_name_scope(name):
     return '/'.join(name.split('/')[:-1])
 
 
-def create_vague_shape_like(shape):
-    make_sure(len(shape) >= 0, "rank should be >= 0")
-    return [-1 for i in enumerate(shape)]
-
-
 def get_temp_directory():
     return os.environ.get("TF2ONNX_TEMP_DIRECTORY", tempfile.mkdtemp())
 
@@ -364,3 +359,71 @@ def save_protobuf(path, message, as_text=False):
     else:
         with open(path, "wb") as f:
             f.write(message.SerializeToString())
+
+
+def is_list_or_tuple(obj):
+    return isinstance(obj, (list, tuple))
+
+
+def is_unknown_dimension(dim):
+    """  Return true if dim is not a positive integer value. """
+    if dim is None or not isinstance(dim, int):
+        return True
+    return dim <= 0
+
+
+def merge_shapes(shape1, shape2):
+    """
+    Merge 2 shapes, return merged shape, choose more specific dimension value from either side.
+    Raise exception for mismatch.
+    """
+    if shape1 is None:
+        return shape2
+    if shape2 is None:
+        return shape1
+
+    make_sure(is_list_or_tuple(shape1), "invalid type for shape1")
+    make_sure(is_list_or_tuple(shape2), "invalid type for shape2")
+    make_sure(len(shape1) == len(shape2), "shapes rank mismatch: shape1=%s, shape2=%s", shape1, shape2)
+
+    merged = []
+    for d1, d2 in zip(shape1, shape2):
+        d = d1
+        if is_unknown_dimension(d1):
+            d = d2
+        elif not is_unknown_dimension(d2):
+            make_sure(d1 == d2, "shapes dimension mismatch: shape1=%s, shape2=%s", shape1, shape2)
+        merged.append(d)
+    return merged
+
+
+def are_shapes_compatible(src, dest):
+    """
+    Returns True iff src is compatible with dest.
+    None is compatible with all shapes, different ranks are not considered as compatible
+    """
+    try:
+        merge_shapes(src, dest)
+        return True
+    except:  # pylint: disable=bare-except
+        return False
+
+
+def are_shapes_equal(src, dest):
+    """ Check whether 2 shapes are equal. """
+    if src is None:
+        return dest is None
+    if dest is None:
+        return src is None
+
+    make_sure(is_list_or_tuple(src), "invalid type for src")
+    make_sure(is_list_or_tuple(dest), "invalid type for dest")
+
+    if len(src) != len(dest):
+        return False
+    return all(i == j for i, j in zip(src, dest))
+
+
+def create_vague_shape_like(shape):
+    make_sure(len(shape) >= 0, "rank should be >= 0")
+    return [-1 for i in enumerate(shape)]