diff --git a/keras2onnx/_builtin.py b/keras2onnx/_builtin.py index 00b2e84d..6d286bf8 100644 --- a/keras2onnx/_builtin.py +++ b/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 diff --git a/keras2onnx/_consts.py b/keras2onnx/_consts.py index 4596f05d..b978262e 100644 --- a/keras2onnx/_consts.py +++ b/keras2onnx/_consts.py @@ -27,6 +27,7 @@ class TYPES: Conv1D = 'Conv1D' Conv2D = 'Conv2D' Cumsum = 'Cumsum' + DepthToSpace = 'DepthToSpace' DepthwiseConv2dNative = 'DepthwiseConv2dNative' Div = 'Div' Einsum = 'Einsum' diff --git a/tests/test_layers.py b/tests/test_layers.py index 336f65ad..62796a4e 100644 --- a/tests/test_layers.py +++ b/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)