Fix TF1-Keras Dilated Conv Export

update patch

skip cpu test for conv3d ncdhw

skip CPU for Conv2D NCHW

update

Signed-off-by: Lei Mao <[email protected]>

Author: leimao

tests/test_backend.py

@@ -1,5 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 """Unit tests using onnx backends."""
 
@@ -727,7 +741,7 @@ def func(x):
                             onnx_feed_dict={_INPUT: x_val_for_onnx})
 
     @skip_tflite("TFlite adds ops that obscure pattern")
-    @check_tf_min_version("2.0")
+    @check_tf_min_version("1.15")
     def test_conv1d_dilations_rewriter(self):
         x_shape = [2, 32, 3]
         x_val = make_xval(x_shape)
@@ -740,7 +754,7 @@ 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")
+    @check_tf_min_version("1.15")
     def test_conv2d_dilations_rewriter(self):
         x_shape = [2, 32, 16, 3]
         x_val = make_xval(x_shape)
@@ -760,7 +774,39 @@ 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")
+    @check_tf_min_version("1.15")
+    @skip_tf_cpu("only tf_gpu can run conv2d with NCHW format")
+    def test_nchw_conv2d_dilations_rewriter(self):
+        x_shape = [2, 3, 32, 16]
+        x_val = make_xval(x_shape)
+        for p in ['SAME', 'VALID']:
+            def func(x):
+                t = tf.keras.layers.Conv2D(
+                    filters=768,
+                    kernel_size=3,
+                    dilation_rate=3,
+                    padding=p,
+                    data_format='channels_first'
+                )
+                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,
+                                graph_validator=lambda g: check_op_count(g, "Reshape", 0, disabled=False))
+            def func(x):
+                t = tf.keras.layers.DepthwiseConv2D(
+                    kernel_size=3,
+                    dilation_rate=3,
+                    padding=p,
+                    data_format='channels_first'
+                )
+                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,
+                                graph_validator=lambda g: check_op_count(g, "Reshape", 0, disabled=False))
+
+    @check_tf_min_version("1.15")
     @skip_tflite("TFlite adds ops that obscure pattern")
     @allow_missing_shapes("Rewriting makes some shapes known")
     def test_conv2d_dilations_rewriter_unknown_shape(self):
@@ -776,7 +822,30 @@ def func():
                             as_session=True, premade_placeholders=True,
                             graph_validator=lambda g: check_op_count(g, "Reshape", 0, disabled=False))
 
-    @check_tf_min_version("2.0")
+    @check_tf_min_version("1.15")
+    @skip_tflite("TFlite adds ops that obscure pattern")
+    @skip_tf_cpu("only tf_gpu can run conv2d with NCHW format")
+    @allow_missing_shapes("Rewriting makes some shapes known")
+    def test_nchw_conv2d_dilations_rewriter_unknown_shape(self):
+        x_shape = [2, 3, 32, 16]
+        x_val = make_xval(x_shape)
+        def func():
+            x = tf_placeholder(tf.float32, [2, 3, None, None], name=_TFINPUT)
+            t = tf.keras.layers.Conv2D(
+                filters=768,
+                kernel_size=3,
+                dilation_rate=3,
+                padding="VALID",
+                data_format='channels_first'
+            )
+            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("1.15")
     def test_conv3d_dilations_rewriter(self):
         x_shape = [2, 32, 16, 8, 3]
         x_val = make_xval(x_shape)
@@ -789,6 +858,26 @@ 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("1.15")
+    @skip_tf_cpu("only tf_gpu can run conv3d with NCDHW format")
+    def test_ncdhw_conv3d_dilations_rewriter(self):
+        x_shape = [2, 3, 32, 16, 8]
+        x_val = make_xval(x_shape)
+        for p in ['SAME', 'VALID']:
+            def func(x):
+                t = tf.keras.layers.Conv3D(
+                    filters=768,
+                    kernel_size=3,
+                    dilation_rate=3,
+                    padding=p,
+                    data_format='channels_first'
+                )
+                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,
+                                graph_validator=lambda g: check_op_count(g, "Reshape", 0, disabled=False))
+
     @skip_tf2("Uses tf.layers")
     def test_conv1d_tf1_dilations_rewriter(self):
         x_shape = [2, 32, 3]

tf2onnx/rewriter/conv_dilations_rewriter.py

@@ -1,5 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 """
 tf2onnx.rewriter.conv_dilations_rewriter - Rewrites the patten used to represent dilations
@@ -13,6 +27,7 @@
 
 
 def rewrite_conv_dilations(g, ops):
+
     pattern1 = \
         OpTypePattern("BatchToSpaceND", name="batch_to_space", inputs=[
             OpTypePattern("DepthwiseConv2dNative|Conv2D|Conv3D", name="conv", inputs=[
@@ -67,14 +82,7 @@ def rewrite_conv_dilations(g, ops):
             if block_shape1 != block_shape2:
                 continue
             ndims = 2 if is_conv_1d else len(block_shape1)
-            data_format = b"NHWC" if ndims == 2 else b"NDHWC"
-            ones = [1] * (ndims + 2)
-            if conv.get_attr_value("dilations", ones) != ones:
-                continue
-            if conv.get_attr_value("strides", ones) != ones:
-                continue
-            if conv.get_attr_value("data_format", data_format) != data_format:
-                continue
+
             if conv.get_attr_value("padding") != b"VALID":
                 continue
 
@@ -114,7 +122,10 @@ def rewrite_conv_dilations(g, ops):
                 g.copy_shape(batch_to_space.output[0], conv.output[0])
                 g.replace_all_inputs(batch_to_space.output[0], conv.output[0])
 
-            conv.set_attr("dilations", [1] + block_shape1 + [1])
+            if conv.get_attr_value("data_format") in [b"NCHW", b"NCDHW"]:
+                conv.set_attr("dilations", [1] + block_shape1)
+            else:
+                conv.set_attr("dilations", [1] + block_shape1 + [1])
             conv.set_attr("padding", pad_mode)
             if pad_mode == "EXPLICIT":
                 conv.set_attr("explicit_paddings", base_pad_flat)