Create a rewriter for fusing Conv2d with BiasAdd

Author: PreethaVeera

tf2onnx/rewriter/__init__.py

@@ -20,6 +20,7 @@
     rewrite_custom_rnn_cell, rewrite_generic_loop
 from tf2onnx.rewriter.thresholded_relu_rewriter import rewrite_thresholded_relu
 from tf2onnx.rewriter.transpose_rewriter import rewrite_transpose
+from tf2onnx.rewriter.conv2d_with_add_rewriter import rewrite_biasadd_with_conv2d
 
 
 __all__ = [
@@ -41,4 +42,5 @@
     "rewrite_bi_direction_gru",
     "rewrite_custom_rnn_cell",
     "rewrite_generic_loop",
+    "rewrite_biasadd_with_conv2d",
 ]

tf2onnx/rewriter/conv2d_with_add_rewriter.py

@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx conv2d op with BiasAdd
+"""
+from tf2onnx import logging
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+logger = logging.getLogger(__name__)
+
+
+# pylint: disable=missing-docstring
+
+def rewrite_biasadd_with_conv2d(g, ops):
+    pattern = \
+        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]]
+
+        # 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

tf2onnx/tfonnx.py

@@ -467,6 +467,7 @@ def compat_handler(ctx, node, **kwargs):
                  rewrite_single_direction_lstm, rewrite_bi_direction_lstm,
                  rewrite_single_direction_gru, rewrite_bi_direction_gru,
                  rewrite_custom_rnn_cell, rewrite_generic_loop, rewrite_cond,
+                 rewrite_biasadd_with_conv2d,
                  ]
 
     if custom_rewriter is not None:

tf2onnx/version.py

@@ -1,3 +1,3 @@
-
-version = '1.6.0'
-git_version = '82f805f8fe7d2fa91e6ca9d39e153712f6887fec'
+
+version = '1.6.0'
+git_version = 'be91553e30582216dfc67e1808e7eb42bb566541'