Add support for tf.nn.depth_to_space lambda

keras2onnx/_builtin.py

@@ -373,6 +373,36 @@ def _calc_explicit_padding(input_size, output_shape, output_padding, kernel_shap
     return pads
 
 
+@converter_func(TYPES.DepthToSpace)
+def convert_tf_depth_to_space(scope, operator, container):
+    node = operator.raw_operator
+    block_size = node.get_attr('block_size')
+    oopb = OnnxOperatorBuilder(container, scope)
+    if _is_nhwc(node):
+        adjusted_input_name = oopb.apply_transpose(operator.input_full_names,
+                                                   name=operator.full_name + '_pre_transpose',
+                                                   perm=[0, 3, 1, 2])
+        depth_to_space_result = oopb.add_node("DepthToSpace",
+                                              adjusted_input_name,
+                                              name=operator.full_name,
+                                              blocksize=node.get_attr('block_size'),
+                                              mode="DCR",
+                                              op_version=11)
+        oopb.apply_op_with_output("apply_transpose",
+                                  depth_to_space_result,
+                                  operator.output_full_names,
+                                  name=operator.full_name + '_post_transpose',
+                                  perm=[0, 2, 3, 1])
+    else:
+        oopb.add_node_with_output("DepthToSpace",
+                                  operator.input_full_names,
+                                  operator.output_full_names,
+                                  name=operator.full_name,
+                                  blocksize=block_size,
+                                  mode="DCR",
+                                  op_version=11)
+
+
 @converter_func(TYPES.DepthwiseConv2dNative)
 def convert_tf_depthwise_conv2d(scope, operator, container):
     node = operator.raw_operator

keras2onnx/_consts.py

@@ -27,6 +27,7 @@ class TYPES:
     Conv1D = 'Conv1D'
     Conv2D = 'Conv2D'
     Cumsum = 'Cumsum'
+    DepthToSpace = 'DepthToSpace'
     DepthwiseConv2dNative = 'DepthwiseConv2dNative'
     Div = 'Div'
     Einsum = 'Einsum'

tests/test_layers.py

@@ -90,6 +90,27 @@ def test_keras_lambda(runner):
     assert runner('onnx_lambda', onnx_model, data, expected)
 
 
+@pytest.mark.skipif(is_tensorflow_older_than('1.12.0'),
+                    reason="tf.nn.depth_to_space not supported.")
+@pytest.mark.skipif(get_maximum_opset_supported() < 11,
+                    reason="DepthToSpace is not supported before opset 11.")
+@pytest.mark.parametrize("data_format", ["NCHW", "NHWC"])
+@pytest.mark.parametrize("input_shape", [(4, 6, 8), (None, None, 8)])
+def test_keras_lambda_depth_to_space(runner, data_format, input_shape):
+    if data_format == "NCHW" and is_tensorflow_older_than("2.1.0"):
+        pytest.skip("tf.nn.depth_to_space with NCHW not supported for Tensorflow older than 2.1.0")
+    model = Sequential()
+    model.add(Lambda(
+        lambda x: tf.nn.depth_to_space(x, block_size=2, data_format=data_format),
+        input_shape=input_shape
+    ))
+
+    onnx_model = keras2onnx.convert_keras(model, 'test_keras_lambda_depth_to_space')
+    data = np.random.rand(3, 4, 6, 8).astype(np.float32)  # batch dimension + 'input_shape'
+    expected = model.predict(data)
+    assert runner('tf_depth_to_space', onnx_model, data, expected)
+
+
 def test_tf_addn(runner):
     input1 = Input(shape=(5, 3, 4), dtype=tf.float32)
     input2 = Input(shape=(5, 3, 4), dtype=tf.float32)