infer shape for tensorflow graph

Author: lucienwang1009

tests/common.py

@@ -43,7 +43,7 @@
 class TestConfig(object):
     def __init__(self):
         self.platform = sys.platform
-        self.tf_version = self._get_tf_version()
+        self.tf_version = utils.get_tf_version()
         self.opset = int(os.environ.get("TF2ONNX_TEST_OPSET", constants.PREFERRED_OPSET))
         self.target = os.environ.get("TF2ONNX_TEST_TARGET", ",".join(constants.DEFAULT_TARGET)).split(',')
         self.backend = os.environ.get("TF2ONNX_TEST_BACKEND", "onnxruntime")
@@ -67,10 +67,6 @@ def is_caffe2_backend(self):
     def is_debug_mode(self):
         return utils.is_debug_mode()
 
-    def _get_tf_version(self):
-        import tensorflow as tf
-        return LooseVersion(tf.__version__)
-
     def _get_backend_version(self):
         version = None
         if self.backend == "onnxruntime":

tests/test_backend.py

@@ -1271,7 +1271,7 @@ def test_randomuniform_int(self):
     def test_randomuniform_dyn_shape(self):
         # test for dynamic shape coming from a shape op
         x_val = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)
-        x = tf.placeholder(x_val.dtype, name=_TFINPUT)
+        x = tf.placeholder(x_val.dtype, [None, 3], name=_TFINPUT)
         x_ = tf.stack([x, x])
         x_ = tf.identity(x_)
         x_ = tf.shape(x_, name="shape")

tests/test_loops.py

@@ -12,14 +12,32 @@
 import tensorflow as tf
 
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main
+from common import unittest_main, check_tf_min_version
 
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
 
 
 class LoopTests(Tf2OnnxBackendTestBase):
 
+    @check_tf_min_version("1.9")
+    def test_simple_while_loop_var_shape(self):
+        # test for while_loop with variant shape variables
+        # may not meet ONNX Loop spec
+        i = tf.placeholder(tf.int32, (1), name="input_1")
+        const = tf.constant(np.array([2], dtype=np.int32))
+
+        c = lambda i: tf.reduce_all(tf.shape(i) < 10)
+        b = lambda i: tf.concat([i, const], 0)
+        r = tf.while_loop(c, b, [i], shape_invariants=[tf.TensorShape([None])])
+
+        _ = tf.identity(r, name="output")
+        input_names_with_port = ["input_1:0"]
+        feed_dict = {"input_1:0": np.array([0], dtype=np.int32)}
+
+        output_names_with_port = ["output:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
+
     def test_simple_while_loop(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         c = lambda i: tf.less(i, 10)

tf2onnx/loader.py

@@ -29,7 +29,7 @@ def freeze_session(sess, keep_var_names=None, output_names=None, clear_devices=T
         freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
         output_names = output_names or []
         output_names += [v.op.name for v in tf.global_variables()]
-        input_graph_def = graph.as_graph_def()
+        input_graph_def = graph.as_graph_def(add_shapes=True)
         if clear_devices:
             for node in input_graph_def.node:
                 node.device = ""

tf2onnx/onnx_opset/controlflow.py

@@ -268,7 +268,7 @@ def version_8(cls, ctx, node, **kwargs):
         make_sure(true_data_shape is not None, "select true data shape cannot be None")
 
         condition_shape = ctx.get_shape(node.input[0])
-        utils.make_sure(condition_shape is not None, "condition shape is None")
+        utils.make_sure(condition_shape is not None, "Shape of {} is None".format(node.input[0]))
         rank = len(condition_shape)
 
         utils.make_sure(rank >= 0, "rank should be >= 0")

tf2onnx/rewriter/loop_rewriter.py

@@ -58,7 +58,9 @@ def rewrite(self, context):
             body_nodes = set(cell_g_info.nodes + cond_g_info.nodes)
             body_outputs = cond_g_info.outputs + cell_g_info.outputs
             for out_tensor_value_info in body_outputs:
-                out_tensor_value_info.shape = utils.create_vague_shape_like(out_tensor_value_info.shape)
+                shape = out_tensor_value_info.shape
+                utils.make_sure(shape is not None, "Shape of {} is None".format(out_tensor_value_info.id))
+                out_tensor_value_info.shape = utils.create_vague_shape_like(shape)
 
             loop_body_g = LoopRewriterBase.construct_graph_from_nodes(self.g, body_nodes, body_outputs)