Make conv dilations rewriter work for dynamic pads (#1470)

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

Co-authored-by: Guenther Schmuelling <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -764,6 +764,22 @@ def func(x):
             self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, rtol=1e-04, atol=1e-2, as_session=True,
                                 graph_validator=lambda g: check_op_count(g, "Reshape", 0, disabled=False))
 
+    @check_tf_min_version("2.0")
+    @skip_tflite("TFlite adds ops that obscure pattern")
+    @allow_missing_shapes("Rewriting makes some shapes known")
+    def test_conv2d_dilations_rewriter_unknown_shape(self):
+        x_shape = [2, 32, 16, 3]
+        x_val = make_xval(x_shape)
+        def func():
+            x = tf_placeholder(tf.float32, [2, None, None, 3], name=_TFINPUT)
+            t = tf.keras.layers.Conv2D(filters=768, kernel_size=3, dilation_rate=3, padding="VALID")
+            t.build(x_shape)
+            y = t.call(x)
+            return tf.identity(y, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, rtol=1e-04, atol=1e-2,
+                            as_session=True, premade_placeholders=True,
+                            graph_validator=lambda g: check_op_count(g, "Reshape", 0, disabled=False))
+
     @check_tf_min_version("2.0")
     def test_conv3d_dilations_rewriter(self):
         x_shape = [2, 32, 16, 8, 3]
@@ -788,6 +804,18 @@ def func(x):
             self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, rtol=1e-04, atol=1e-2, as_session=True,
                                 graph_validator=lambda g: check_op_count(g, "Reshape", 0, disabled=False))
 
+    @skip_tf2("Uses tf.layers")
+    def test_conv1d_tf1_dilations_rewriter_unknown_shape(self):
+        x_shape = [2, 32, 3]
+        x_val = make_xval(x_shape)
+        def func():
+            x = tf_placeholder(tf.float32, [2, None, 3], name=_TFINPUT)
+            y = tf.layers.conv1d(x, filters=768, kernel_size=3, dilation_rate=3, padding="VALID", name="conv1")
+            return tf.identity(y, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, rtol=1e-04, atol=1e-2,
+                            as_session=True, premade_placeholders=True,
+                            graph_validator=lambda g: check_op_count(g, "Reshape", 0, disabled=False))
+
     def test_lrn_default(self):
         x_shape = [1, 3, 4, 3]
         x_val = np.arange(1, 1 + np.prod(x_shape)).astype("float32").reshape(x_shape)

tf2onnx/rewriter/conv_dilations_rewriter.py

@@ -6,6 +6,7 @@
 pat = SpaceToBatchND->DepthwiseConv2dNative->BatchToSpaceND
 """
 
+import numpy as np
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 
 # pylint: disable=invalid-name,unused-argument,missing-docstring, unused-variable
@@ -18,30 +19,30 @@ def rewrite_conv_dilations(g, ops):
                 OpTypePattern("SpaceToBatchND", name="space_to_batch", inputs=[
                     OpTypePattern("*"),
                     OpTypePattern("Const|ConstV2"),
-                    OpTypePattern("Const|ConstV2"),
+                    OpTypePattern("*"),
                 ]),
                 OpTypePattern("*"),
             ]),
             OpTypePattern("Const|ConstV2"),
-            OpTypePattern("Const|ConstV2"),
+            OpTypePattern("*"),
         ])
     pattern2 = \
         OpTypePattern("BatchToSpaceND", name="batch_to_space", inputs=[
             OpTypePattern("Squeeze", name="squeeze", inputs=[
-                OpTypePattern("DepthwiseConv2dNative|Conv2D|Conv3D", name="conv", inputs=[
+                OpTypePattern("Conv2D", name="conv", inputs=[
                     OpTypePattern("ExpandDims", name="expand", inputs=[
                         OpTypePattern("SpaceToBatchND", name="space_to_batch", inputs=[
                             OpTypePattern("*"),
                             OpTypePattern("Const|ConstV2"),
-                            OpTypePattern("Const|ConstV2"),
+                            OpTypePattern("*"),
                         ]),
                         OpTypePattern("Const|ConstV2"),
                     ]),
                     OpTypePattern("*"),
                 ]),
             ]),
             OpTypePattern("Const|ConstV2"),
-            OpTypePattern("Const|ConstV2"),
+            OpTypePattern("*"),
         ])
 
     for pattern in [pattern1, pattern2]:
@@ -61,9 +62,7 @@ def rewrite_conv_dilations(g, ops):
                     continue
 
             block_shape1 = space_to_batch.inputs[1].get_tensor_value(as_list=True)
-            paddings = space_to_batch.inputs[2].get_tensor_value(as_list=True)
             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
@@ -79,15 +78,27 @@ def rewrite_conv_dilations(g, ops):
             if conv.get_attr_value("padding") != b"VALID":
                 continue
 
-
-            base_start_pad = [p[0] for p in paddings]
-            if any(c[0] != 0 for c in crops):
-                continue
-            base_end_pad = [p[1] - c[1] for p, c in zip(paddings, crops)]
-            if not all(0 <= p[1] - bp < bs for p, bp, bs in zip(paddings, base_end_pad, block_shape1)):
-                continue
+            if space_to_batch.inputs[2].is_const() and batch_to_space.inputs[2].is_const():
+                paddings = space_to_batch.inputs[2].get_tensor_value(as_list=True)
+                crops = batch_to_space.inputs[2].get_tensor_value(as_list=True)
+                base_start_pad = [p[0] for p in paddings]
+                if any(c[0] != 0 for c in crops):
+                    continue
+                base_end_pad = [p[1] - c[1] for p, c in zip(paddings, crops)]
+                if not all(0 <= p[1] - bp < bs for p, bp, bs in zip(paddings, base_end_pad, block_shape1)):
+                    continue
+                pad_mode = "EXPLICIT"
+                if is_conv_1d:
+                    base_start_pad = [0] + base_start_pad
+                    base_end_pad = [0] + base_end_pad
+                base_pad_flat = [0, 0] + [x for s, e in zip(base_start_pad, base_end_pad) for x in [s, e]] + [0, 0]
+            else:
+                pad_mode = determine_pad_mode(space_to_batch.inputs[2])
+                if pad_mode is None:
+                    continue
 
             if is_conv_1d:
+                block_shape1 = [1] + block_shape1
                 inp = space_to_batch.input[0]
                 g.replace_inputs(expand, [inp, expand.input[1]])
                 g.copy_shape(batch_to_space.output[0], squeeze.output[0])
@@ -96,20 +107,39 @@ def rewrite_conv_dilations(g, ops):
                 g.set_shape(squeeze.input[0], squeeze_out_shape[:1] + [1] + squeeze_out_shape[1:])
                 expand_inp_shape = g.get_shape(expand.input[0])
                 g.set_shape(expand.output[0], expand_inp_shape[:1] + [1] + expand_inp_shape[1:])
-
-                base_start_pad = [0] + base_start_pad
-                base_end_pad = [0] + base_end_pad
-                block_shape1 = [1] + block_shape1
             else:
                 inp = space_to_batch.input[0]
                 kernel = conv.input[1]
                 g.replace_inputs(conv, [inp, kernel])
                 g.copy_shape(batch_to_space.output[0], conv.output[0])
                 g.replace_all_inputs(batch_to_space.output[0], conv.output[0])
 
-            base_pad_flat = [0, 0] + [x for s, e in zip(base_start_pad, base_end_pad) for x in [s, e]] + [0, 0]
             conv.set_attr("dilations", [1] + block_shape1 + [1])
-            conv.set_attr("explicit_paddings", base_pad_flat)
-            conv.set_attr("padding", "EXPLICIT")
+            conv.set_attr("padding", pad_mode)
+            if pad_mode == "EXPLICIT":
+                conv.set_attr("explicit_paddings", base_pad_flat)
 
     return g.get_nodes()
+
+def determine_pad_mode(paddings_inp_node):
+    tensor_ops = set(["Concat", "ConcatV2", "ConcatV3", "StridedSlice", "Pack", "ExpandDims", "Identity"])
+    while paddings_inp_node.type in tensor_ops:
+        non_const = [inp for inp in paddings_inp_node.inputs if not inp.is_const()]
+        if len(non_const) == 0:
+            return None
+        paddings_inp_node = non_const[0]
+    if paddings_inp_node.type == "FloorMod":
+        return "VALID"
+    if paddings_inp_node.type in ["Add", "AddV2"]:
+        if paddings_inp_node.inputs[0].type == "FloorMod":
+            pad = paddings_inp_node.inputs[1]
+        elif paddings_inp_node.inputs[1].type == "FloorMod":
+            pad = paddings_inp_node.inputs[0]
+        else:
+            return None
+        if pad.is_const():
+            if np.any(pad.get_tensor_value(as_list=False)):
+                #return "SAME"   ORT doesn't implement dilations for SAME autopadding yet
+                return None
+            return "VALID"
+    return None