Implement fused tfjs ops (#1620)

TomWildenhain-Microsoft · web-flow · commit 8c081efde1bb · 2021-07-15T20:17:10.000-07:00
* Squashed commit of the following: commit de562af Author: Tom Wildenhain <tomwi@microsoft.com> Date: Tue Jul 13 18:33:56 2021 -0700 Pylint commit 39dc6d2 Author: Tom Wildenhain <tomwi@microsoft.com> Date: Tue Jul 13 18:25:38 2021 -0700 Rename tfjs_runner.py Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> commit cdcccfc Author: Tom Wildenhain <tomwi@microsoft.com> Date: Tue Jul 13 18:13:08 2021 -0700 Enable tfjs unit tests Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> commit b6efae4 Merge: de5ac2c c5aba9a Author: Tom Wildenhain <tomwi@microsoft.com> Date: Mon Jul 12 22:00:14 2021 -0700 Merge branch 'master' into tom/tfjs2 commit de5ac2c Author: Tom Wildenhain <tomwi@microsoft.com> Date: Mon Jul 12 21:59:17 2021 -0700 Exclude tfjs tests with broken tf->tfjs conversions Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> commit d92d7a3 Author: Tom Wildenhain <tomwi@microsoft.com> Date: Mon Jul 12 13:55:44 2021 -0700 TFJS fixes Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> commit 103761f Merge: e4889a7 b65b05e Author: Tom Wildenhain <tomwi@microsoft.com> Date: Fri Jul 9 13:34:39 2021 -0700 Merge branch 'master' into tom/tfjs2 commit e4889a7 Merge: a0c2ee4 664430c Author: Tom Wildenhain <tomwi@microsoft.com> Date: Thu Jul 8 21:17:29 2021 -0700 Merge branch 'master' into tom/tfjs2 commit a0c2ee4 Author: Tom Wildenhain <tomwi@microsoft.com> Date: Thu Jul 8 21:17:09 2021 -0700 Improve testing for tfjs Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> commit ca0f9bd Author: Tom Wildenhain <tomwi@microsoft.com> Date: Thu Jul 8 20:41:26 2021 -0700 Add tests for tfjs Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> commit 91ffbca Author: Tom Wildenhain <tomwi@microsoft.com> Date: Wed Jul 7 20:50:15 2021 -0700 Implement tfjs conversion Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> commit 5190778 Merge: e04b40f 448d61d Author: Tom Wildenhain <tomwi@microsoft.com> Date: Wed Jul 7 14:44:08 2021 -0700 Merge branch 'master' into tom/tfjs2 commit e04b40f Merge: 58ed277 dc1d0e1 Author: Tom Wildenhain <tomwi@microsoft.com> Date: Tue Jul 6 20:40:21 2021 -0700 Merge branch 'tom/refactor_tf2onnx2' into tom/tfjs2 commit dc1d0e1 Author: Tom Wildenhain <tomwi@microsoft.com> Date: Tue Jul 6 19:44:20 2021 -0700 Refactor tflite/tf logic out of tfonnx into tf_utils and tflite_utils Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> commit 58ed277 Merge: b4c3186 552b0a2 Author: Tom Wildenhain <tomwi@microsoft.com> Date: Tue Jul 6 18:56:02 2021 -0700 Merge branch 'master' into tom/tfjs2 commit b4c3186 Author: Tom Wildenhain <tomwi@microsoft.com> Date: Fri Jul 2 18:31:49 2021 -0700 WIP Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> * reorganize Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> * pylint Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> * Improve documentation Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> * Pylint Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> * Implement conversion of _Fused ops seen in tfjs Signed-off-by: Tom Wildenhain <tomwi@microsoft.com> * Enable tests with fused ops Signed-off-by: Tom Wildenhain <tomwi@microsoft.com>
diff --git a/tests/test_backend.py b/tests/test_backend.py
@@ -2770,7 +2770,6 @@ def func(x, mean, offset, var):
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: mean_val, _INPUT2: offset_val, _INPUT3: var_val})
 
     @check_opset_min_version(7, "batchnorm")
-    @skip_tfjs("Unsupported _FusedConv2D op")   # TODO: implement this
     def test_conv2d_batchnorm_fusion(self):
         x_shape = [1, 28, 28, 2]
         x_val = np.random.random_sample(x_shape).astype(np.float32)
@@ -4272,7 +4271,6 @@ def func(a, b, c):
                             graph_validator=lambda g: check_op_count(g, "Gemm", 1))
 
     # test for gemm pattern4: A*B + C [addbias] - 1D bias!
-    @skip_tfjs("Unsupported _FusedMatMul op")   # TODO: implement this
     def test_gemm_pattern4(self):
         max_number = 10
         m = np.random.randint(max_number)
diff --git a/tf2onnx/rewriter/__init__.py b/tf2onnx/rewriter/__init__.py
@@ -22,6 +22,7 @@
 from tf2onnx.rewriter.layer_normalization_rewriter import rewrite_layer_normalization
 from tf2onnx.rewriter.ragged_variant_shape_rewriter import rewrite_ragged_variant_shape
 from tf2onnx.rewriter.lstm_tf2_rewriter import rewriter_lstm_tf2
+from tf2onnx.rewriter.fused_op_rewriter import rewrite_fused_ops
 
 
 __all__ = [
@@ -48,5 +49,6 @@
     "rewrite_layer_normalization",
     "rewrite_conv_dilations",
     "rewrite_ragged_variant_shape",
-    "rewriter_lstm_tf2"
+    "rewriter_lstm_tf2",
+    "rewrite_fused_ops",
 ]
diff --git a/tf2onnx/rewriter/fused_op_rewriter.py b/tf2onnx/rewriter/fused_op_rewriter.py
@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""
+tf2onnx.rewriter.fused_op_rewriter - rewrite tensorflow _Fused ops from grappler into other tf ops
+"""
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_fused_ops(g, ops):
+    for node in ops:
+        if node.type in ["_FusedConv2D", "_FusedMatMul"]:
+            op_types = [op.decode() for op in node.get_attr_value("fused_ops")]
+            extra_inputs = node.input[2:]
+            g.replace_inputs(node, node.input[:2])
+            last_output = node.output[0]
+            node.type = node.type.replace("_Fused", "")
+            dtype = g.get_dtype(node.output[0])
+            shape = g.get_shape(node.output[0])
+            first_node = None
+            for op in op_types:
+                new_node = g.make_node(op, [last_output] + extra_inputs, skip_conversion=False,
+                                       op_name_scope=node.name, dtypes=[dtype], shapes=[shape])
+                last_output = new_node.output[0]
+                if first_node is None:
+                    first_node = new_node
+                extra_inputs = []
+
+            consumers = [n for n in g.find_output_consumers(node.output[0]) if n != first_node]
+            g.replace_all_inputs(node.output[0], last_output, consumers)
+
+    return g.get_nodes()
diff --git a/tf2onnx/tfonnx.py b/tf2onnx/tfonnx.py
@@ -554,6 +554,7 @@ def compat_handler(ctx, node, **kwargs):
         # single directional
         rewrite_constant_fold,
         rewrite_quantize_and_dequantize,
+        rewrite_fused_ops,
         rewrite_transpose,
         rewrite_flatten,
         rewrite_random_uniform,
