Merge pull request #56 from onnx/gs/onnx-1.2

support for tf.unstack

Author: guschmue

tests/test_backend.py

@@ -834,18 +834,30 @@ def test_topk(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
-    def test_stack_axis0(self):
-        x_val = [np.random.randn(3, 4).astype("float32") for _ in range(10)]
-        x = [tf.constant(x_val[i], dtype=tf.float32) for i in range(10)]
-        x_ = tf.stack(x, axis=0)
-        output = tf.identity(x_, name=_TFOUTPUT)
-        actual, expected = self._run(output, {}, {})
-        self.assertAllClose(expected, actual)
-
-    def test_stack_axis1(self):
-        x_val = [np.random.randn(3, 4).astype("float32") for _ in range(10)]
-        x = [tf.constant(x_val[i], dtype=tf.float32) for i in range(10)]
-        x_ = tf.stack(x, axis=1)
+    def test_stack_axis(self):
+        for axis in [0, 1]:
+            tf.reset_default_graph()
+            x_val = [np.random.randn(3, 4).astype("float32") for _ in range(10)]
+            x = [tf.constant(x_val[i], dtype=tf.float32) for i in range(10)]
+            x_ = tf.stack(x, axis=axis)
+            output = tf.identity(x_, name=_TFOUTPUT)
+            actual, expected = self._run(output, {}, {})
+            self.assertAllClose(expected, actual)
+
+    def test_unstack_axis(self):
+        for axis in [0, 1]:
+            tf.reset_default_graph()
+            x_val = np.random.randn(10, 3, 4).astype("float32")
+            x = tf.constant(x_val, dtype=tf.float32)
+            x_ = tf.unstack(x, axis=axis)
+            output = tf.identity(x_, name=_TFOUTPUT)
+            actual, expected = self._run(output, {}, {})
+            self.assertAllClose(expected, actual)
+
+    def test_unstack_axis1(self):
+        x_val = np.random.randn(10, 3, 4).astype("float32")
+        x = tf.constant(x_val, dtype=tf.float32)
+        x_ = tf.unstack(x, axis=1)
         output = tf.identity(x_, name=_TFOUTPUT)
         actual, expected = self._run(output, {}, {})
         self.assertAllClose(expected, actual)

tf2onnx/tfonnx.py

@@ -853,6 +853,23 @@ def pack_op(ctx, node, name, args):
     ctx.replace_all_inputs(ctx.get_nodes(), node.output[0], output_name)
     return [concat] + nodes
 
+def unpack_op(ctx, node, name, args):
+    # hack to make up for the missing onnx unpack op
+    axis = node.get_attr("axis").i
+    # split the tensor into n outputs
+    node.type = "Split"
+    nodes = [node]
+    # for each output we need to squeeze axis
+    for i, n in enumerate(node.output):
+        op_name = utils.make_name(node.name)
+        output_name = op_name + ":" + str(i)
+        new_node = Node(helper.make_node("Squeeze", [n], [output_name], name=op_name, axes=[axis]), ctx)
+        nodes.append(new_node)
+        ctx.copy_shape(n, output_name)
+        ctx.replace_all_inputs(ctx.get_nodes(), n, output_name)
+    return nodes
+
+
 # pylint: enable=W0613,C0111,W0612
 
 # map tensorflow ops to onnx ops. The format below is
@@ -939,6 +956,7 @@ def pack_op(ctx, node, name, args):
     "TopKV2": (topk_op, []),
     "SpaceToDepth": (spacetodepth_op, []),
     "Pack": (pack_op, []),
+    "Unpack": (unpack_op, []),
 }
 
 _OPSET_5 = {