Merge pull request #615 from zhijxu-MS/fix_graph_matcher

Fix bug and refactor graph matcher

Author: wayuanho

tests/test_backend.py

@@ -17,6 +17,7 @@
 # pylint reports unused-wildcard-import which is false positive, __all__ is defined in common
 from common import *  # pylint: disable=wildcard-import,unused-wildcard-import
 from tf2onnx import constants
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument
 
@@ -87,6 +88,7 @@ def get_conv_getdata(kind=1):
     else:
         raise ValueError("kind not known")
 
+
 def get_maxpoolwithargmax_getdata():
     data = [
         ('SAME', [1, 3, 3, 1], [1, 3, 3, 1], [1, 2, 2, 1]),
@@ -99,6 +101,7 @@ def get_maxpoolwithargmax_getdata():
     for idx, v in enumerate(data):
         yield (idx,) + v
 
+
 class BackendTests(Tf2OnnxBackendTestBase):
     def _run_test_case(self, output_names_with_port, feed_dict, **kwargs):
         kwargs["convert_var_to_const"] = False
@@ -2014,7 +2017,6 @@ def test_reverse_sequence_time_major(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-
     @check_opset_min_version(10, "ReverseSequence")
     def test_reversev2_constant_axis(self):
         # Tests for constant axis.
@@ -2034,7 +2036,6 @@ def test_reversev2_constant_axis(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-
     @check_opset_min_version(10, "ReverseSequence")
     def test_reversev2_vector_axis(self):
         x_val_shape = [1, 2, 3, 4]
@@ -2060,7 +2061,6 @@ def test_reversev2_vector_axis(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-
     @check_opset_min_version(10, "ReverseSequence")
     def test_reversev2_1D_tensor(self):
         # For tensors with 1 dimension and no axis to reverse.
@@ -2072,7 +2072,6 @@ def test_reversev2_1D_tensor(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-
     @check_opset_min_version(8, "where")
     def test_where(self):
         x_val = np.array([1, 2, -3, 4, -5, -6, -7, 8, 9, 0], dtype=np.float32)
@@ -2552,5 +2551,36 @@ def test_selu(self):
         _ = tf.identity(y, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    def test_graph_matcher(self):
+        shape = [2, 6]
+        x_val = np.random.random(shape).astype(np.float32)
+        y_val = np.random.random(shape).astype(np.float32)
+        z_val = np.random.random(shape).astype(np.float32)
+        x = tf.placeholder(tf.float32, shape, name=_TFINPUT)
+        y = tf.placeholder(tf.float32, shape, name=_TFINPUT1)
+        z = tf.placeholder(tf.float32, shape, name=_TFINPUT2)
+        tmp1 = x + y
+        tmp2 = x - y
+        tmp3 = tf.multiply(tmp1, z)
+        tmp4 = tf.multiply(tmp2, z)
+        _ = tf.add(tmp4, tmp3, name=_TFOUTPUT)
+        onnx_graph = self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: y_val, _INPUT2: z_val})
+        pattern = \
+            OpTypePattern('Add', name='output', inputs=[
+                OpTypePattern('Mul', inputs=[
+                    OpTypePattern('Add', name='input1'),
+                    OpTypePattern('*', name='input2')]),
+                OpTypePattern('Mul', inputs=[
+                    OpTypePattern('Sub', name='input1'),
+                    OpTypePattern('*', name='input2')])])
+
+        matcher = GraphMatcher(pattern, allow_reorder=False)
+        match_results = list(matcher.match_ops(onnx_graph.get_nodes()))
+        self.assertTrue(len(match_results) == 0)
+        matcher = GraphMatcher(pattern, allow_reorder=True)
+        match_results = list(matcher.match_ops(onnx_graph.get_nodes()))
+        self.assertTrue(len(match_results) == 1)
+
+
 if __name__ == '__main__':
     unittest_main()

tf2onnx/graph_matcher.py

@@ -19,12 +19,10 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import copy
-
+from itertools import permutations
 import six
 
 
-
 class OpTypePattern(object):
     """A tree pattern that matches TF expressions with certain op types."""
 
@@ -137,6 +135,16 @@ def __init__(self, pattern, allow_reorder=False):
         self._pattern = pattern
         self._allow_reorder = allow_reorder
 
+    @staticmethod
+    def _is_op_type_same(op, pattern):
+        if pattern.op_type == "*":
+            return True
+
+        if op.type in pattern.op_type.split('|'):
+            return True
+
+        return False
+
     def _match_pattern(self, pattern, op, tensor):
         """Returns whether an TF expression rooted at `op` matches `pattern`.
 
@@ -151,48 +159,55 @@ def _match_pattern(self, pattern, op, tensor):
             pattern tree.
 
         Returns:
-          True if an TF expression rooted at `op` matches `pattern`.
+          if matched return True and match_list whose elem is [pattern, op, tensor]
+          else return False
+        the condition that op is matched with pattern:
+        1 op is same:
+          if pattern.op_type is None or *, then treat as same
+          or op.type in pattern.op_type.split("|")
+        2 op.inputs are same with pattern.inputs:
+          if not pattern.inputs, then treat as same
+          otherwise, iteratively compare input nodes with pattern.
         """
+        match_list = []
         if pattern.op_type is None:
-            return True
-
-        if pattern.op_type != '*':
-            if op is None or op.type not in pattern.op_type.split('|'):
-                return False
+            return True, match_list
 
-        self._match_result.add(pattern, op, tensor)
-        # print("matched", ",".join([op.type + "|" + op.name for op in self._match_result.get_nodes()]))
+        if self._is_op_type_same(op, pattern):
+            match_list.append([pattern, op, tensor])
+        else:
+            return False, match_list
 
         if not pattern.inputs:
             # If pattern.inputs is empty, skips the rest and accepts all the inputs.
-            return True
+            return True, match_list
 
         if not op or len(op.inputs) != len(pattern.inputs):
-            return False
+            return False, match_list
 
         if self._allow_reorder:
-            inputs = [None] * len(op.inputs)
-            wanted = copy.copy(pattern.inputs)
-            for idx, i in enumerate(op.inputs):
-                for j in range(len(wanted)):  # pylint: disable=consider-using-enumerate
-                    if i.type == wanted[j].op_type:
-                        inputs[idx] = wanted[j]
-                        del wanted[j]
-                        break
-            for idx, i in enumerate(inputs):
-                if i is None:
-                    inputs[idx] = wanted[0]
-                    del wanted[0]
-            pat = list(zip(op.inputs, inputs))
+            pattern_inputs_list = permutations(pattern.inputs)
         else:
-            pat = list(zip(op.inputs, pattern.inputs))
-
-        ret = []
-        for input_tensor, input_pattern in pat:
-            # print("MATCHING", input_pattern.op_type, input_tensor.type)
-            r = self._match_pattern(input_pattern, input_tensor, input_tensor)
-            ret.append(r)
-        return all(ret)
+            pattern_inputs_list = [pattern.inputs]
+
+        for possible_pattern_inputs in pattern_inputs_list:
+            pat = list(zip(op.inputs, possible_pattern_inputs))
+            match_flag_of_inputs = []
+            match_lists_of_inputs = []
+            for input_tensor, input_pattern in pat:
+                # print("MATCHING", input_pattern.op_type, input_tensor.type)
+                flag, match_list_of_input = self._match_pattern(input_pattern, input_tensor, input_tensor)
+                match_flag_of_inputs.append(flag)
+                match_lists_of_inputs.extend(match_list_of_input)
+
+            if all(match_flag_of_inputs):
+                match_list.extend(match_lists_of_inputs)
+                return True, match_list
+        return False, match_list
+
+    def _parse_match_list_to_match_result(self, match_list):
+        for pattern, op, tensor in match_list:
+            self._match_result.add(pattern, op, tensor)
 
     def match_op(self, op):
         """Matches `op` against `self._pattern`.
@@ -205,8 +220,10 @@ def match_op(self, op):
           None.
         """
         self._match_result = MatchResult()
-        if not self._match_pattern(self._pattern, op, tensor=None):
+        match_flag, match_list = self._match_pattern(self._pattern, op, tensor=None)
+        if not match_flag:
             return None
+        self._parse_match_list_to_match_result(match_list)
         return self._match_result
 
     def match_ops(self, ops):

tf2onnx/rewriter/lstm_rewriter.py

@@ -125,10 +125,14 @@ def _get_weight_and_bias_for_lstm_cell(self, context):
     def parse_attributes(self, context):
         if self.lstm_cell_type == RNNUnitType.LSTMBlockCell:
             lstm_block_cell = context.cell_match.get_op("lstm_block_cell")
-            clip = float(lstm_block_cell.get_attr("cell_clip").f)
+            clip = lstm_block_cell.get_attr_value("cell_clip")
             # current LSTM op cannot handle clip
             if clip > 0:
                 return False
+
+            use_peephole = lstm_block_cell.get_attr_value("use_peephole")
+            if use_peephole:
+                return False
         return True
 
     def _ct_variable_finder(self, context):