Merge pull request #316 from zhijxu-MS/map_fn

add test cases of map_fn, fix bug of shape_inference

Author: nbcsm

tests/test_backend.py

@@ -211,7 +211,6 @@ def test_multinomial(self):
         self._run_test_case([_OUTPUT], {_INPUT: x_val}, check_value=False,
                             check_shape=True, check_dtype=True)
 
-
     @unittest.skipIf(BACKEND in ["caffe2"], "not supported correctly in caffe2")
     @unittest.skipIf(*support_op_conversion_since(7, "multinomial"))
     def test_multinomial1(self):

tests/test_loops.py

@@ -16,6 +16,7 @@
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
 
+
 class LoopTests(Tf2OnnxBackendTestBase):
 
     def test_simple_while_loop(self):
@@ -31,7 +32,6 @@ def test_simple_while_loop(self):
         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_2(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         c = lambda i: tf.logical_and(tf.less(i, 10), tf.greater_equal(i, 3))
@@ -45,7 +45,6 @@ def test_simple_while_loop_2(self):
         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_while_loop_with_ta_write(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         output_ta = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True)
@@ -68,7 +67,6 @@ def b(i, out_ta):
         output_names_with_port = ["output:0", "i:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
-
     def test_while_loop_with_ta_read(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         inputs = tf.placeholder(tf.float32, (10,), name="input_2")
@@ -82,6 +80,7 @@ def test_while_loop_with_ta_read(self):
         c = lambda i, *_: tf.logical_and(tf.less(i, 10), i >= 0)
         res = tf.constant(0.)
         res2 = tf.constant(1.)
+
         def b(i, res, res2):
             new_i = tf.add(i, 1)
             x = input_ta.read(i)
@@ -113,6 +112,7 @@ def test_while_loop_with_ta_read_reference_outer_input_directly(self):
         c = lambda i, *_: tf.logical_and(tf.less(i, 10), i >= 0)
         res = tf.constant(0.)
         res2 = tf.constant(1.)
+
         def b(i, res, res2):
             new_i = tf.add(i, 1)
             x = input_ta.read(i)
@@ -132,7 +132,6 @@ def b(i, res, res2):
         output_names_with_port = ["i:0", "x:0", "y:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
-
     def test_while_loop_with_ta_read_and_write(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         inputs = tf.placeholder(tf.float32, (10,), name="input_2")
@@ -160,5 +159,42 @@ def b(i, out_ta):
         output_names_with_port = ["i:0", "output_ta:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
+    def test_map_fn(self):
+        def fn0(elem):
+            res = elem + elem * elem
+            return res
+
+        def fn1(elem):
+            res = elem[0] * elem[1] + elem[0]
+            return res
+
+        x_val = 100 * np.random.random_sample([2, 10]).astype(np.float32)
+        y_val = 100 * np.random.random_sample([2, 10]).astype(np.float32)
+
+        # test fn0
+        x = tf.placeholder(tf.float32, shape=x_val.shape, name="input_0")
+        x_ = tf.identity(x)
+        res_ = tf.map_fn(fn0, x_, dtype=tf.float32)
+        _ = tf.identity(res_, name="output_0")
+        feed_dict = {"input_0:0": x_val}
+        input_names_with_port = ["input_0:0"]
+        output_names_with_port = ["output_0:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-5)
+        tf.reset_default_graph()
+
+        # test fn1
+        x = tf.placeholder(tf.float32, shape=x_val.shape, name="input_0")
+        y = tf.placeholder(tf.float32, shape=y_val.shape, name="input_1")
+        x_ = tf.identity(x)
+        y_ = tf.identity(y)
+        res_ = tf.map_fn(fn1, (x_, y_), dtype=tf.float32)
+        _ = tf.identity(res_, name="output_0")
+        feed_dict = {"input_0:0": x_val, "input_1:0": y_val}
+        input_names_with_port = ["input_0:0", "input_1:0"]
+        output_names_with_port = ["output_0:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-5)
+        tf.reset_default_graph()
+
+
 if __name__ == '__main__':
     Tf2OnnxBackendTestBase.trigger(LoopTests)

tf2onnx/shape_inference.py

@@ -10,6 +10,7 @@
 from __future__ import unicode_literals
 import logging
 from onnx import onnx_pb
+from tf2onnx import utils
 
 # pylint: disable=logging-not-lazy,missing-docstring,consider-swap-variables
 
@@ -91,10 +92,10 @@ def infer_shape_for_node(g, node):
         shape_attr = node.get_attr("shape")
         new_shape = None
         if shape_attr.type == onnx_pb.TensorProto.INT32:
-            new_shape = shape_attr.ints
+            new_shape = list(shape_attr.ints)
         elif shape_attr.type == onnx_pb.TensorProto.FLOAT:
             # for scalar placeholder, it's type is float
-            val = shape_attr.floats
+            val = list(shape_attr.floats)
             if val:
                 raise ValueError("placeholder shape has floats value, and not scalar value")
             else:
@@ -243,13 +244,24 @@ def infer_output_shapes_with_partial_inputs(g, node):
         return True
 
     if node.type == "TensorArrayGatherV3":
-        # TensorArrayGatherV3's output: all of the elements in the TensorArray,
-        # concatenated along a new axis (the new dimension 0)
-        flow_in_node = node.inputs[2]
-        if flow_in_node.type != "Exit":
+        # TensorArrayGatherV3's output: all of the elem in the TensorArray,
+        # concatenated along a new axis (the new dimension 0), so shape of TensorArray should be found first.
+        # And TensorArrayWrite will write elem to TensorArray, so shape of TensorArray can be got from TensorArrayWrite
+        # so the process is: first find TensorArrayWrite and then get TensorArray's shape,
+        # and finally add one dim to the shape is shape of TensorArrayGather
+
+        handle_node = node.inputs[0]
+        if handle_node.type != "TensorArrayV3":
             return False
 
-        shape = g.get_shape(flow_in_node.output[0])
+        # find TensorArrayWrite
+        tensor_array_write_node = _find_tensorarray_write(g, handle_node)
+        if not tensor_array_write_node:
+            return False
+        # get TensorArray shape from input tensor of the found TensorArrayWrite node
+        value_node = tensor_array_write_node.inputs[2]
+        shape = g.get_shape(value_node.output[0])
+        # update TensorArray's shape info
         if shape is not None:
             new_shape = [-1] + shape
             g.set_shape(node.output[0], new_shape)
@@ -350,3 +362,16 @@ def broadcast_shape_inference(shape_0, shape_1):
             return None
         i -= 1
     return new_shape
+
+
+def _find_tensorarray_write(graph, node):
+    utils.make_sure(node.type == "TensorArrayV3", "node should be tensorarray")
+
+    tensor_array_consumers = graph.find_output_consumers(node.output[0])
+    for i in tensor_array_consumers:
+        if i.type == "Enter":
+            consumer_nodes = graph.find_output_consumers(i.output[0])
+            for j in consumer_nodes:
+                if j.type == "TensorArrayWriteV3":
+                    return j
+    return None