Optimize SpaceToBatchND sequence to DepthwiseConv2dNative dilations (#1461)

* Add support for Explicit padding

Signed-off-by: Tom Wildenhain <[email protected]>

* Optimize SpaceToBatchND sequence to DepthwiseConv2dNative dilations

Signed-off-by: Tom Wildenhain <[email protected]>

* pylint

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -3411,6 +3411,18 @@ def func(x, z):
                     return batch_to_space_nd(x, y, z, name=_TFOUTPUT)
                 self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT2: z_val})
 
+    def test_depthwise_dilations_pattern(self):
+        x_val = np.random.random_sample([1, 33, 34, 960]).astype(np.float32)
+        kernel = np.random.random_sample([3, 3, 960, 1]).astype(np.float32)
+        block_size = np.array([3, 3], np.int64)
+        pad = np.array([[2, 4], [5, 3]])
+        crop = np.array([[0, 0], [0, 0]])
+        def func(x):
+            y = space_to_batch_nd(x, block_size, pad)
+            z = tf.nn.depthwise_conv2d(y, kernel, strides=[1, 1, 1, 1], padding='VALID')
+            return batch_to_space_nd(z, block_size, crop, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
     @check_opset_min_version(11, "SpaceToBatchND")
     def test_space_to_batchnd_non_const_7d(self):
         x_type, y_type, z_type = np.float32, np.int64, np.int64

tf2onnx/rewriter/__init__.py

@@ -8,6 +8,7 @@
 
 from tf2onnx.rewriter.cond_rewriter import rewrite_cond
 from tf2onnx.rewriter.conv2d_with_pad_rewriter import rewrite_conv2d_with_pad
+from tf2onnx.rewriter.depthwise_conv_dilations_rewriter import rewrite_depthwise_conv_dilations
 from tf2onnx.rewriter.dropout_rewriter import rewrite_dropout
 from tf2onnx.rewriter.eye_rewriter import rewrite_eye
 from tf2onnx.rewriter.flatten_rewriter import rewrite_flatten
@@ -46,5 +47,6 @@
     "rewrite_generic_loop",
     "rewrite_biasadd_with_conv2d",
     "rewrite_quantize_and_dequantize",
-    "rewrite_layer_normalization"
+    "rewrite_layer_normalization",
+    "rewrite_depthwise_conv_dilations"
 ]

tf2onnx/rewriter/depthwise_conv_dilations_rewriter.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""
+tf2onnx.rewriter.depthwise_conv_dilations_rewriter - Rewrites the patten used to represent dilations
+pat = SpaceToBatchND->DepthwiseConv2dNative->BatchToSpaceND
+"""
+
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+# pylint: disable=invalid-name,unused-argument,missing-docstring, unused-variable
+
+
+def rewrite_depthwise_conv_dilations(g, ops):
+    pattern1 = \
+        OpTypePattern("BatchToSpaceND", name="batch_to_space", inputs=[
+            OpTypePattern("DepthwiseConv2dNative", name="depthwise_conv", inputs=[
+                OpTypePattern("SpaceToBatchND", name="space_to_batch", inputs=[
+                    OpTypePattern("*"),
+                    OpTypePattern("Const|ConstV2"),
+                    OpTypePattern("Const|ConstV2"),
+                ]),
+                OpTypePattern("*"),
+            ]),
+            OpTypePattern("Const|ConstV2"),
+            OpTypePattern("Const|ConstV2"),
+        ])
+
+    for pattern in [pattern1]:
+        matcher = GraphMatcher(pattern, allow_reorder=False)
+        match_results = list(matcher.match_ops(ops))
+        for match_result in match_results:
+            space_to_batch = match_result.get_op("space_to_batch")
+            depthwise_conv = match_result.get_op("depthwise_conv")
+            batch_to_space = match_result.get_op("batch_to_space")
+
+            block_shape1 = space_to_batch.inputs[1].get_tensor_value(as_list=True)
+            paddings = space_to_batch.inputs[2].get_tensor_value(as_list=False).flatten().tolist()
+            block_shape2 = batch_to_space.inputs[1].get_tensor_value(as_list=True)
+            crops = batch_to_space.inputs[2].get_tensor_value(as_list=True)
+            if block_shape1 != block_shape2:
+                continue
+            if depthwise_conv.get_attr_value("dilations", [1, 1, 1, 1]) != [1, 1, 1, 1]:
+                continue
+            if depthwise_conv.get_attr_value("strides", [1, 1, 1, 1]) != [1, 1, 1, 1]:
+                continue
+            if depthwise_conv.get_attr_value("data_format", b"NHWC") != b"NHWC":
+                continue
+            if depthwise_conv.get_attr_value("padding") != b"VALID":
+                continue
+            if crops != [[0, 0], [0, 0]]:
+                continue
+
+            inp = space_to_batch.input[0]
+            kernel = depthwise_conv.input[1]
+
+            g.replace_inputs(depthwise_conv, [inp, kernel])
+            depthwise_conv.set_attr("dilations", [1] + block_shape1 + [1])
+            depthwise_conv.set_attr("explicit_paddings", [0, 0] + paddings + [0, 0])
+            depthwise_conv.set_attr("padding", "EXPLICIT")
+            g.copy_shape(batch_to_space.output[0], depthwise_conv.output[0])
+            g.replace_all_inputs(batch_to_space.output[0], depthwise_conv.output[0])
+
+    return g.get_nodes()

tf2onnx/tfonnx.py

@@ -604,6 +604,7 @@ def compat_handler(ctx, node, **kwargs):
         rewrite_random_uniform_fold_const,
         rewrite_random_normal,
         rewrite_dropout,
+        rewrite_depthwise_conv_dilations,
         rewrite_eye,
         rewrite_leakyrelu,
         rewrite_thresholded_relu,