Implement rewriter for LSTM nodes in tf2 (#1584)

TomWildenhain-Microsoft · web-flow · commit 02061e990319 · 2021-06-30T18:49:45.000Z
Signed-off-by: Tom Wildenhain &lt;tomwi@microsoft.com&gt;
diff --git a/tests/test_lstm.py b/tests/test_lstm.py
@@ -9,7 +9,7 @@
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import variable_scope
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main, check_opset_after_tf_version, skip_tf2, skip_tf_versions
+from common import unittest_main, check_opset_after_tf_version, skip_tf2, skip_tf_versions, check_op_count
 
 from tf2onnx.tf_loader import is_tf2
 
@@ -36,12 +36,22 @@
 
 class LSTMTests(Tf2OnnxBackendTestBase):
 
-    def run_test_case(self, *args, **kwargs):  #pylint: disable=arguments-differ
+    def run_test_case(self, *args, require_lstm_count=1, **kwargs):  #pylint: disable=arguments-differ
         # TF LSTM has an unknown dim
         tmp = self.config.allow_missing_shapes
         self.config.allow_missing_shapes = True
+        def graph_validator(g):
+            good = True
+            if "graph_validator" in kwargs:
+                good = good and kwargs["graph_validator"](g)
+            if require_lstm_count is None or ":" not in g.outputs[0]:
+                # Skip checks for tflite graphs (no ":" in outputs)
+                return good
+            good = good and check_op_count(g, "LSTM", require_lstm_count, disabled=False)
+            good = good and check_op_count(g, "Loop", 0, disabled=False)
+            return good
         try:
-            super().run_test_case(*args, **kwargs)
+            super().run_test_case(*args, graph_validator=graph_validator, **kwargs)
         finally:
             self.config.allow_missing_shapes = tmp
 
@@ -385,7 +395,8 @@ def func(x):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0", "cell_state:0"]
-        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
+        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06,
+                           require_lstm_count=2)
 
     @check_opset_after_tf_version("1.15", 8, "might need Scan")
     @skip_tf2()     # Still failing likely due to inconsistent random number initialization
diff --git a/tf2onnx/graph.py b/tf2onnx/graph.py
@@ -464,6 +464,8 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
         # A list of index, output tuples of potential scan outputs in this graph
         # Used by the tflite while loop handler
         self.scan_outputs = []
+        # Used by lstm_tf2_rewriter to indicate this subgraph is an LSTM cell
+        self.lstm_rewriter_context = None
         self.func_inputs = []
         self.ragged_variant_list_reads = []
         self.ragged_variant_list_writes = []
diff --git a/tf2onnx/rewriter/__init__.py b/tf2onnx/rewriter/__init__.py
@@ -21,6 +21,7 @@
 from tf2onnx.rewriter.quantization_ops_rewriter import rewrite_quantize_and_dequantize
 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
 
 
 __all__ = [
@@ -46,5 +47,6 @@
     "rewrite_quantize_and_dequantize",
     "rewrite_layer_normalization",
     "rewrite_conv_dilations",
-    "rewrite_ragged_variant_shape"
+    "rewrite_ragged_variant_shape",
+    "rewriter_lstm_tf2"
 ]
diff --git a/tf2onnx/rewriter/lstm_tf2_rewriter.py b/tf2onnx/rewriter/lstm_tf2_rewriter.py
@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""
+tf2onnx.rewriter.lstm_tf2_rewriter - Rewrites LSTM pattern used by tf2.
+"""
+
+import numpy as np
+from tf2onnx.graph_matcher import GraphMatcher
+from tf2onnx.rewriter.rnn_utils import make_lstmcell_pattern
+from tf2onnx.tf_loader import find_function
+from tf2onnx.rewriter.lstm_rewriter_base import LSTMContext
+from tf2onnx.rewriter.lstm_rewriter import LSTMRewriter
+from tf2onnx.graph_builder import GraphBuilder
+from tf2onnx import utils
+
+# pylint: disable=invalid-name,unused-argument,missing-docstring, unused-variable
+
+
+def rewriter_lstm_tf2(g, ops):
+    pattern1 = make_lstmcell_pattern("Identity")
+
+    for pattern in [pattern1]:
+        matcher = GraphMatcher(pattern, allow_reorder=False)
+        match_results = list(matcher.match_ops(ops))
+        for match_result in match_results:
+            concat = match_result.get_op("xh")
+            if len(concat.inputs) != 3:
+                continue
+            get_item = concat.inputs[0]
+            if not get_item.type == "TensorListGetItem":
+                continue
+            x_e = get_item.inputs[0]
+            if not x_e.is_graph_input():
+                continue
+            x_idx = g.input_names.index(x_e.output[0])
+
+            ht_mul = match_result.get_op("ht")
+            final_consumers = g.find_output_consumers(ht_mul.output[0])
+            select_ops = [n for n in final_consumers if n.type == "Select"]
+            def has_tensor_list_consumer(n):
+                return any(c.type == "TensorListSetItem" for c in g.find_output_consumers(n.output[0]))
+            select_ops = [n for n in select_ops if has_tensor_list_consumer(n)]
+            if len(select_ops) == 1:
+                greater_eq = select_ops[0].inputs[0]
+                if greater_eq.type != "GreaterEqual":
+                    continue
+                seq_len = greater_eq.inputs[1]
+                if not seq_len.is_graph_input():
+                    continue
+                seq_len_idx = g.input_names.index(seq_len.output[0])
+                final_consumers = g.find_output_consumers(select_ops[0].output[0])
+            else:
+                seq_len_idx = None
+
+            tensor_set_items = [n for n in final_consumers if n.type == "TensorListSetItem"]
+            if len(tensor_set_items) != 1:
+                continue
+
+            if not tensor_set_items[0].inputs[0].is_graph_input():
+                continue
+            out_idx = g.input_names.index(tensor_set_items[0].input[0])
+
+            if concat.inputs[1].is_graph_input():
+                # c and h are separate
+                h_idx = g.input_names.index(concat.input[1])
+                c_e = match_result.get_op("c")
+                if not c_e.is_graph_input():
+                    continue
+                c_idx = g.input_names.index(c_e.output[0])
+                ch_info = {
+                    "state_is_tuple": True,
+                    "c_idx": c_idx,
+                    "h_idx": h_idx,
+                }
+            else:
+                # c and h are concatenated
+                if not concat.inputs[1].type == "Slice":
+                    continue
+                ch_e = concat.inputs[1].inputs[0]
+                if not ch_e.is_graph_input():
+                    continue
+                ch_idx = g.input_names.index(ch_e.output[0])
+
+                c_e = match_result.get_op("c")
+                if not c_e.type == "Slice" or c_e.input[0] != ch_e.output[0]:
+                    continue
+                ch_info = {
+                    "state_is_tuple": False,
+                    "ch_idx": ch_idx,
+                }
+
+            w_e = match_result.get_op("cell_kernel")
+            if not w_e.is_graph_input():
+                continue
+            w_idx = g.input_names.index(w_e.output[0])
+
+            bias_add = match_result.get_op("bias_add")
+            if bias_add is not None and bias_add.data_format != "NHWC":
+                continue
+
+            b_e = match_result.get_op("cell_bias")
+            if not b_e.is_graph_input():
+                continue
+            b_idx = g.input_names.index(b_e.output[0])
+
+            ft_bias_node = match_result.get_op("ft_bias")
+            if not ft_bias_node.is_const():
+                continue
+            if g.get_dtype(ft_bias_node.output[0]) != g.get_dtype(b_e.output[0]):
+                continue
+            ft_bias = ft_bias_node.get_tensor_value(as_list=False)
+
+            g.lstm_rewriter_context = {
+                "x_idx": x_idx,
+                "out_idx": out_idx,
+                "weight_idx": w_idx,
+                "bias_idx": b_idx,
+                "ft_bias": ft_bias,
+                "seq_len_idx": seq_len_idx,
+                **ch_info
+            }
+
+    for op in ops:
+        if op.is_while():
+            body_graph = find_function(op.get_attr_str("body"))
+            if body_graph.lstm_rewriter_context is None:
+                continue
+            body_context = body_graph.lstm_rewriter_context
+            w = op.input[body_context["weight_idx"]]
+            b = op.input[body_context["bias_idx"]]
+            if not g.is_const(w) or not g.is_const(b):
+                continue
+            w_const = g.get_tensor_value(w, as_list=False)
+            b_const = g.get_tensor_value(b, as_list=False)
+
+            if body_context["state_is_tuple"]:
+                initial_c_sq = op.input[body_context["c_idx"]]
+                initial_h_sq = op.input[body_context["h_idx"]]
+                initial_c = GraphBuilder(g).make_unsqueeze({"data": initial_c_sq, "axes": [0]})
+                initial_h = GraphBuilder(g).make_unsqueeze({"data": initial_h_sq, "axes": [0]})
+            else:
+                initial_ch = op.input[body_context["ch_idx"]]
+                if not g.is_const(initial_ch):
+                    continue
+                initial_ch_const = g.get_tensor_value(initial_ch, as_list=False)
+                if not len(initial_ch_const.shape) == 2:
+                    continue
+                initial_ch_const = np.expand_dims(initial_ch_const, axis=0)
+                initial_c_const, initial_h_const = np.split(initial_ch_const, 2, axis=2)
+                initial_c = g.make_const(utils.make_name("initial_c"), initial_c_const).output[0]
+                initial_h = g.make_const(utils.make_name("initial_h"), initial_h_const).output[0]
+
+            context = LSTMContext()
+            context.weights.append({"weight": w_const, "bias": b_const, "ft_bias": body_context["ft_bias"]})
+            context.onnx_input_ids.append({})
+            context.input_size.append(None)
+            context.hidden_size.append(None)
+            context.attributes.append({})
+            tensor_array_inp = op.inputs[body_context["x_idx"]]
+            if not tensor_array_inp.type == "TensorListFromTensor":
+                continue
+
+            final_consumers = g.find_output_consumers(op.output[body_context["out_idx"]])
+            output_ys = [n.output[0] for n in final_consumers if n.type == "TensorListStack"]
+
+            context.onnx_input_ids[0]["X"] = tensor_array_inp.input[0]
+            if body_context["seq_len_idx"] is None:
+                context.onnx_input_ids[0]["sequence_lens"] = ""
+            else:
+                context.onnx_input_ids[0]["sequence_lens"] = op.input[body_context["seq_len_idx"]]
+            context.onnx_input_ids[0]["initial_c"] = initial_c
+            context.onnx_input_ids[0]["initial_h"] = initial_h
+
+            lstm_rewriter = LSTMRewriter(g)
+            lstm_rewriter.num_lstm_layers = 1
+            lstm_rewriter.process_weights_and_bias(context)
+            lstm_node = lstm_rewriter.create_rnn_node(context)[0]
+            squeeze_output = GraphBuilder(g).make_squeeze({"data": lstm_node.output[0], "axes": [1]})
+            for output in output_ys:
+                g.replace_all_inputs(output, squeeze_output)
+
+            if body_context["state_is_tuple"]:
+                c_squeeze = GraphBuilder(g).make_squeeze({"data": lstm_node.output[2], "axes": [0]})
+                h_squeeze = GraphBuilder(g).make_squeeze({"data": lstm_node.output[1], "axes": [0]})
+                g.replace_all_inputs(op.output[body_context["c_idx"]], c_squeeze)
+                g.replace_all_inputs(op.output[body_context["h_idx"]], h_squeeze)
+            else:
+                concat_ch = g.make_node("Concat", [lstm_node.output[2], lstm_node.output[1]],
+                                        attr={"axis": 2}).output[0]
+                ch_squeeze = GraphBuilder(g).make_squeeze({"data": concat_ch, "axes": [0]})
+                ch_output = op.output[body_context["ch_idx"]]
+                g.replace_all_inputs(ch_output, ch_squeeze)
+
+    return g.get_nodes()
diff --git a/tf2onnx/rewriter/rnn_utils.py b/tf2onnx/rewriter/rnn_utils.py
@@ -30,44 +30,51 @@ class REWRITER_RESULT(Enum):
 
 # TensorFlow LSTMCell/BasicLSTMCell computation graph matching
 
-xc_pattern = \
-    OpTypePattern('Split', inputs=[
+_make_xc_pattern_memo = {}
+
+def make_xc_pattern(enter_or_id="Enter"):
+    return OpTypePattern('Split', inputs=[
         OpTypePattern("Const"), # axis for split
         OpTypePattern("BiasAdd", name="bias_add", inputs=[
             OpTypePattern("MatMul", inputs=[
                 OpTypePattern("ConcatV2|Concat", name="xh"),
-                OpTypePattern("Enter", inputs=[
+                OpTypePattern(enter_or_id, inputs=[
                     OpTypePattern("*", name="cell_kernel"),
                 ]),
             ]),
-            OpTypePattern("Enter", inputs=[
+            OpTypePattern(enter_or_id, inputs=[
                 OpTypePattern("*", name="cell_bias"),
             ]),
         ]),
     ])
 
-lstmcell_pattern = \
-    OpTypePattern('Mul', name='ht', inputs=[
-        OpTypePattern("Sigmoid", name="ot", inputs=[xc_pattern]),
+xc_pattern = make_xc_pattern()
+
+def make_lstmcell_pattern(enter_or_id="Enter"):
+    my_xc_pattern = make_xc_pattern(enter_or_id)
+    return OpTypePattern('Mul', name='ht', inputs=[
+        OpTypePattern("Sigmoid", name="ot", inputs=[my_xc_pattern]),
         OpTypePattern('Tanh', inputs=[
             OpTypePattern("Add|AddV2", name="ct", inputs=[
                 OpTypePattern("Mul", name="ct_identity_consumer", inputs=[
                     OpTypePattern("Sigmoid", name="ft", inputs=[
                         OpTypePattern("Add|AddV2", inputs=[
-                            xc_pattern,
+                            my_xc_pattern,
                             OpTypePattern("*", name="ft_bias"),
                         ]),
                     ]),
-                    OpTypePattern("*"),
+                    OpTypePattern("*", name="c"),
                 ]),
                 OpTypePattern("Mul", inputs=[
-                    OpTypePattern("Sigmoid", name="it", inputs=[xc_pattern]),
-                    OpTypePattern("Tanh", name="gt", inputs=[xc_pattern]),
+                    OpTypePattern("Sigmoid", name="it", inputs=[my_xc_pattern]),
+                    OpTypePattern("Tanh", name="gt", inputs=[my_xc_pattern]),
                 ]),
             ]),
         ]),
     ])
 
+lstmcell_pattern = make_lstmcell_pattern()
+
 xc_pattern_optimized = \
     OpTypePattern('Split', inputs=[
         OpTypePattern("Const"),
diff --git a/tf2onnx/tfonnx.py b/tf2onnx/tfonnx.py
@@ -606,6 +606,7 @@ def compat_handler(ctx, node, **kwargs):
         rewrite_leakyrelu,
         rewrite_thresholded_relu,
         rewrite_conv2d_with_pad,
+        rewriter_lstm_tf2,
         rewrite_single_direction_lstm,
         # bi-directional
         rewrite_bi_direction_lstm,
