fix Conv2D Bias Add fuse

Signed-off-by: hwangdeyu <[email protected]>

Author: hwangdeyu

tests/keras2onnx_unit_tests/test_layers.py

@@ -1198,6 +1198,28 @@ def test_conv2d_transpose_2(runner, padding):
     assert runner(onnx_model.graph.name, onnx_model, data, expected)
 
 
+def test_conv2d_bias_add(runner):
+    size = 128
+    input_img = Input((size, size, 64))
+    x = tf.keras.layers.ZeroPadding2D(
+        padding=(0, 4),
+        data_format="channels_first",
+        name="padding")(input_img)
+    y = tf.keras.layers.Conv2D(
+        filters=1,
+        kernel_size=(9, 9),
+        strides=(1, 1),
+        use_bias=True,
+        data_format="channels_first",
+        name="conv2d",
+    )(x)
+    model = tf.keras.Model(inputs=input_img, outputs=y)
+    data = np.random.rand(1, size, size, 64).astype(np.float32)
+    onnx_model = convert_keras(model, model.name)
+    expected = model.predict(data)
+    assert runner(onnx_model.graph.name, onnx_model, data, expected)
+
+
 def test_conv2d_padding_same(conv2_runner):
     conv2_runner(3, 5, (2, 2), (1, 1), (5, 5), padding='same')
     conv2_runner(8, 16, (1, 1), (2, 2), (60, 60), padding='same')

tf2onnx/rewriter/__init__.py

@@ -28,7 +28,6 @@
 
 __all__ = [
     "rewrite_cond",
-    "rewrite_conv2d_with_pad",
     "rewrite_dropout",
     "rewrite_eye",
     "rewrite_flatten",
@@ -49,6 +48,7 @@
     "rewrite_quantize_and_dequantize",
     "rewrite_layer_normalization",
     "rewrite_conv_dilations",
+    "rewrite_conv2d_with_pad",
     "rewrite_ragged_variant_shape",
     "rewriter_lstm_tf2",
     "rewrite_gru_tf2",

tf2onnx/rewriter/conv2d_with_add_rewriter.py

@@ -13,31 +13,40 @@
 # pylint: disable=missing-docstring
 
 def rewrite_biasadd_with_conv2d(g, ops):
-    pattern = \
+    pattern1 = \
         OpTypePattern('BiasAdd', name='biasadd', inputs=[
             OpTypePattern('Conv2D|Conv2DBackpropInput', name='conv', inputs=['*', '*']), '*'])
-    matcher = GraphMatcher(pattern)
-    match_results = list(matcher.match_ops(ops))
-    for match in match_results:
-        biasadd = match.get_op('biasadd')
-        conv = match.get_op('conv')
-
-        #backup the conv and biasadd values
-        conv_type = conv.type
-        conv_input = conv.input
-        conv_attr = conv.attr
-        dtype = g.get_dtype(conv.output[0])
-        shape = g.get_shape(conv.output[0])
-        conv_name = biasadd.name
-        conv_output = biasadd.output
-        conv_inputs = [conv_input[0], conv_input[1], biasadd.input[1]]
-
-        if len(g.find_output_consumers(conv.output[0])) > 1:
-            continue
-        # Remove the Conv and BiasAdd node
-        g.remove_node(conv.name)
-        g.remove_node(biasadd.name)
-
-        g.make_node(conv_type, conv_inputs, attr=conv_attr, name=conv_name, outputs=conv_output,
-                    shapes=[shape], dtypes=[dtype], skip_conversion=False)
+    pattern2 = \
+        OpTypePattern('BiasAdd', name='biasadd', inputs=[
+            OpTypePattern('Conv2D|Conv2DBackpropInput', name='conv', inputs=[
+                '*', '*', '*']), '*'], allow_reorder=True)
+
+    for pattern in [pattern1, pattern2]:
+        matcher = GraphMatcher(pattern)
+        match_results = list(matcher.match_ops(ops))
+        for match in match_results:
+            biasadd = match.get_op('biasadd')
+            conv = match.get_op('conv')
+
+            # Backup the conv and biasadd values
+            conv_type = conv.type
+            conv_input = conv.input
+            conv_attr = conv.attr
+            dtype = g.get_dtype(conv.output[0])
+            shape = g.get_shape(conv.output[0])
+            conv_name = biasadd.name
+            conv_output = biasadd.output
+            if pattern == pattern2:
+                conv_inputs = [conv_input[0], conv_input[1], conv_input[2], biasadd.input[1]]
+            else:
+                conv_inputs = [conv_input[0], conv_input[1], biasadd.input[1]]
+
+            if len(g.find_output_consumers(conv.output[0])) > 1:
+                continue
+            # Remove the Conv and BiasAdd node
+            g.remove_node(conv.name)
+            g.remove_node(biasadd.name)
+
+            g.make_node(conv_type, conv_inputs, attr=conv_attr, name=conv_name, outputs=conv_output,
+                        shapes=[shape], dtypes=[dtype], skip_conversion=False)
     return ops