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

add support for tf.space_to_depth()

Author: guschmue

tests/test_backend.py

@@ -773,6 +773,16 @@ def test_topk(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
+    @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "Space2Depth not implemented, works on onnxmsrtnext")
+    def test_space_to_depth(self):
+        x_val = make_xval([1, 2, 2, 1])
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        x_ = tf.space_to_depth(x, block_size=2)
+        output  = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
+
     @unittest.skipIf(OPSET < 6, "supported since opset 6")
     def test_addn(self):
         x_val = np.arange(3*2*3).astype("float32")

tf2onnx/tfonnx.py

@@ -793,6 +793,13 @@ def tile_op7(ctx, node, name, args):
     return _convert_shapenode_to_int64(ctx, node, 1)
 
 
+def spacetodepth_op(ctx, node, name, args):
+    block_size = node.get_attr("block_size")
+    node.set_attr("blocksize", block_size.i)
+    nodes = conv_convert_inputs(ctx, node, with_kernel=False)
+    return nodes
+
+
 # pylint: enable=W0613,C0111,W0612
 
 # map tensorflow ops to onnx ops. The format below is
@@ -876,6 +883,7 @@ def tile_op7(ctx, node, name, args):
     "Tanh": (direct_op, []),
     "Transpose": (transpose_op, []),
     "TopKV2": (topk_op, []),
+    "SpaceToDepth": (spacetodepth_op, []),
 }
 
 _OPSET_5 = {

tf2onnx/utils.py

@@ -75,7 +75,8 @@
     'min', 'seed', 'ends', 'paddings', 'to', 'gamma', 'width_scale', 'normalize_variance', 'group', 'ratio', 'values',
     'dtype', 'output_shape', 'spatial', 'split', 'input_forget', 'keepdims', 'transA', 'auto_pad', 'border', 'low',
     'linear_before_reset', 'height_scale', 'output_padding', 'shape', 'kernel_shape', 'epsilon', 'size', 'starts',
-    'direction', 'max', 'clip', 'across_channels', 'value', 'strides', 'extra_shape', 'scales', 'k', 'sample_size'
+    'direction', 'max', 'clip', 'across_channels', 'value', 'strides', 'extra_shape', 'scales', 'k', 'sample_size',
+    'blocksize'
 }