add gemm_rewriter and the corresponding test (#597)

JiayingGaoo · wayuanho · commit dbd3569fd989 · 2019-07-01T11:14:36.000+08:00
* add gemm_rewriter and the corresponding tests

* add safe_remove_nodes function in gemm_rewriter.py

* Coding improvements
diff --git a/tests/test_backend.py b/tests/test_backend.py
@@ -2571,6 +2571,80 @@ def test_selu(self):
         _ = tf.identity(y, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    # test for gemm pattern0: alpha*A*B + beta*C
+    def test_gemm_pattern0(self):
+        max_number = 10
+        m = np.random.randint(max_number)
+        n = np.random.randint(max_number)
+        k = np.random.randint(max_number)
+        x_val1 = np.random.rand(m, n).astype("float32")
+        x_val2 = np.random.rand(n, k).astype("float32")
+        x_val3 = np.random.rand(m, k).astype("float32")
+        a = tf.placeholder(tf.float32, x_val1.shape, name=_TFINPUT)
+        b = tf.placeholder(tf.float32, x_val2.shape, name=_TFINPUT1)
+        c = tf.placeholder(tf.float32, x_val3.shape, name=_TFINPUT2)
+        alpha = tf.constant(1.0, dtype=tf.float32)
+        beta = tf.constant(2.0, dtype=tf.float32)
+        mul1 = tf.multiply(alpha, tf.matmul(a, b))
+        mul2 = tf.multiply(beta, c)
+        x_ = mul1 + mul2
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2, _INPUT2: x_val3},
+                            graph_validator=lambda g: check_op_count(g, "Gemm", 1))
+
+    # test for gemm pattern1: alpha*A*B + C
+    def test_gemm_pattern1(self):
+        max_number = 10
+        m = np.random.randint(max_number)
+        n = np.random.randint(max_number)
+        k = np.random.randint(max_number)
+        x_val1 = np.random.rand(m, n).astype("float32")
+        x_val2 = np.random.rand(n, k).astype("float32")
+        x_val3 = np.random.rand(m, k).astype("float32")
+        a = tf.placeholder(tf.float32, x_val1.shape, name=_TFINPUT)
+        b = tf.placeholder(tf.float32, x_val2.shape, name=_TFINPUT1)
+        c = tf.placeholder(tf.float32, x_val3.shape, name=_TFINPUT2)
+        alpha = tf.constant(1.0, dtype=tf.float32)
+        x_ = tf.multiply(alpha, tf.matmul(a, b)) + c
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2, _INPUT2: x_val3},
+                            graph_validator=lambda g: check_op_count(g, "Gemm", 1))
+
+    # test for gemm pattern2: A*B + beta*C
+    def test_gemm_pattern2(self):
+        max_number = 10
+        m = np.random.randint(max_number)
+        n = np.random.randint(max_number)
+        k = np.random.randint(max_number)
+        x_val1 = np.random.rand(m, n).astype("float32")
+        x_val2 = np.random.rand(n, k).astype("float32")
+        x_val3 = np.random.rand(m, k).astype("float32")
+        a = tf.placeholder(tf.float32, x_val1.shape, name=_TFINPUT)
+        b = tf.placeholder(tf.float32, x_val2.shape, name=_TFINPUT1)
+        c = tf.placeholder(tf.float32, x_val3.shape, name=_TFINPUT2)
+        beta = tf.constant(2.0, dtype=tf.float32)
+        x_ = tf.matmul(a, b) + tf.multiply(beta, c)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2, _INPUT2: x_val3},
+                            graph_validator=lambda g: check_op_count(g, "Gemm", 1))
+
+    # test for gemm pattern3: A*B + C
+    def test_gemm_pattern3(self):
+        max_number = 10
+        m = np.random.randint(max_number)
+        n = np.random.randint(max_number)
+        k = np.random.randint(max_number)
+        x_val1 = np.random.rand(m, n).astype("float32")
+        x_val2 = np.random.rand(n, k).astype("float32")
+        x_val3 = np.random.rand(m, k).astype("float32")
+        a = tf.placeholder(tf.float32, x_val1.shape, name=_TFINPUT)
+        b = tf.placeholder(tf.float32, x_val2.shape, name=_TFINPUT1)
+        c = tf.placeholder(tf.float32, x_val3.shape, name=_TFINPUT2)
+        x_ = tf.matmul(a, b) + c
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2, _INPUT2: x_val3},
+                            graph_validator=lambda g: check_op_count(g, "Gemm", 1))
+
     def test_graph_matcher(self):
         shape = [2, 6]
         x_val = np.random.random(shape).astype(np.float32)
diff --git a/tf2onnx/graph_matcher.py b/tf2onnx/graph_matcher.py
@@ -107,16 +107,30 @@ def _to_pattern(self, pattern_or_name):
             return pattern_or_name
 
         if isinstance(pattern_or_name, six.text_type):
-            return self._name_to_pattern[pattern_or_name]
+            return self._name_to_pattern.get(pattern_or_name)
 
         raise ValueError('pattern_or_name has type %s. Expect OpTypePattern or str.'
                          % type(pattern_or_name))
 
-    def get_op(self, pattern_or_name):
-        return self._pattern_to_op_tensor[self._to_pattern(pattern_or_name)][0]
+    def get_op(self, pattern_or_name, default=None):
+        """
+        For now, if the op can not be effectively obtained, then the function will return the default
+        instead of an error.
+        """
+        op_and_tensor = self._pattern_to_op_tensor.get(self._to_pattern(pattern_or_name))
+        if op_and_tensor:
+            return op_and_tensor[0]
+        return default
 
-    def get_tensor(self, pattern_or_name):
-        return self._pattern_to_op_tensor[self._to_pattern(pattern_or_name)][1]
+    def get_tensor(self, pattern_or_name, default=None):
+        """
+        For now, if the tensor can not be effectively obtained, then the function will return the default
+        instead of an error.
+        """
+        op_and_tensor = self._pattern_to_op_tensor.get(self._to_pattern(pattern_or_name))
+        if op_and_tensor:
+            return op_and_tensor[1]
+        return default
 
     def get_nodes(self):
         return [n[0] for n in self._pattern_to_op_tensor.values()]
diff --git a/tf2onnx/rewriter/__init__.py b/tf2onnx/rewriter/__init__.py
@@ -9,6 +9,7 @@
 from tf2onnx.rewriter.cond_rewriter import rewrite_cond
 from tf2onnx.rewriter.random_uniform import rewrite_random_uniform, rewrite_random_uniform_fold_const
 from tf2onnx.rewriter.leakyrelu_rewriter import rewrite_leakyrelu
+from tf2onnx.rewriter.gemm_rewriter import rewrite_gemm
 from tf2onnx.rewriter.eye_rewriter import rewrite_eye
 from tf2onnx.rewriter.thresholded_relu_rewriter import rewrite_thresholded_relu
 from tf2onnx.rewriter.rnn import rewrite_single_direction_lstm, rewrite_bi_direction_lstm, \
@@ -27,5 +28,6 @@
     "rewrite_single_direction_gru",
     "rewrite_bi_direction_gru",
     "rewrite_custom_rnn_cell",
+    "rewrite_gemm",
     "rewrite_generic_loop"
 ]
diff --git a/tf2onnx/rewriter/gemm_rewriter.py b/tf2onnx/rewriter/gemm_rewriter.py
@@ -0,0 +1,102 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewrite - rewrite tensorflow subgraph to onnx gemm op
+"""
+import logging
+from onnx import onnx_pb
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+# pylint: disable=missing-docstring
+
+def rewrite_gemm(g, ops):
+    if g.opset <= 6:
+        return ops
+
+    # pattern0: alpha*A*B + beta*C
+    pattern0 = \
+        OpTypePattern('Add', name='add', inputs=[
+            OpTypePattern('Mul', name='mul1', inputs=[
+                OpTypePattern('Const', name='alpha'),
+                OpTypePattern('MatMul', name='matmul')
+            ]),
+            OpTypePattern('Mul', name='mul2', inputs=[
+                OpTypePattern('Const', name='beta'),
+                OpTypePattern('*', name='C')
+            ])
+        ])
+
+    # pattern1: alpha*A*B + C
+    pattern1 = \
+        OpTypePattern('Add', name='add', inputs=[
+            OpTypePattern('Mul', name='mul1', inputs=[
+                OpTypePattern('MatMul', name='matmul'),
+                OpTypePattern('Const', name='alpha')
+            ]),
+            OpTypePattern('*', name='C'),
+        ])
+
+    # pattern2: A*B + beta*C
+    pattern2 = \
+        OpTypePattern('Add', name='add', inputs=[
+            OpTypePattern('MatMul', name='matmul'),
+            OpTypePattern('Mul', name='mul2', inputs=[
+                OpTypePattern('Const', name='beta'),
+                OpTypePattern('*', name='C')
+            ])
+        ])
+
+    # pattern3: A*B + C
+    pattern3 = \
+        OpTypePattern('Add', name='add', inputs=[
+            OpTypePattern('MatMul', name='matmul'),
+            OpTypePattern('*', name='C'),
+        ])
+
+    pattern_list = [pattern0, pattern1, pattern2, pattern3]
+
+    for pattern in pattern_list:
+        matcher = GraphMatcher(pattern, allow_reorder=True)
+        match_results = list(matcher.match_ops(ops))
+        if match_results:
+            for match in match_results:
+                matmul_node = match.get_op("matmul")
+
+                if g.get_dtype(matmul_node.input[0]) != onnx_pb.TensorProto.FLOAT:
+                    logging.warning(u"For now, onnxruntime only support float32 type for Gemm rewriter")
+                    continue
+
+                attr, is_valid = get_gemm_attr(match)
+                if not is_valid:
+                    continue
+
+                add_node = match.get_op('add')
+                input_c_node = match.get_op("C")
+                a_edge_name = matmul_node.input[0]
+                b_edge_name = matmul_node.input[1]
+                c_edge_name = input_c_node.output[0]
+
+                gemm = g.make_node("Gemm", inputs=[a_edge_name, b_edge_name, c_edge_name],
+                                   attr=attr,
+                                   shapes=[g.get_shape(add_node.output[0])],
+                                   dtypes=[g.get_dtype(add_node.output[0])])
+
+                ops.append(gemm)
+                g.replace_all_inputs(ops, add_node.output[0], gemm.output[0])
+                to_delete = [add_node, matmul_node]
+                g.safe_remove_nodes(to_delete)
+    return ops
+
+def get_gemm_attr(match):
+    attr = {}
+    for arg in ["alpha", "beta"]:
+        arg_op = match.get_op(arg)
+        if arg_op is not None:
+            match_args = arg_op.get_tensor_value()
+            if isinstance(match_args, list):
+                if len(match_args) != 1:
+                    return attr, False
+                match_args = match_args[0]
+            attr[arg] = match_args
+    return attr, True
diff --git a/tf2onnx/tfonnx.py b/tf2onnx/tfonnx.py
@@ -768,13 +768,13 @@ def compat_handler(ctx, node, **kwargs):
 
     # pre-processing graph rewrites
     # bi-directional re-writer should be placed after single directional re-writer
-    rewriters = [rewrite_transpose, rewrite_flatten,
+    rewriters = [rewrite_transpose, rewrite_flatten, rewrite_gemm,
                  rewrite_random_uniform, rewrite_random_uniform_fold_const,
                  rewrite_random_normal, rewrite_dropout, rewrite_eye,
                  rewrite_leakyrelu, rewrite_thresholded_relu, rewrite_conv2d_with_pad,
                  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_custom_rnn_cell, rewrite_generic_loop, rewrite_cond,
                  ]
 
     if custom_rewriter is not None:
