Merge pull request #532 from lucienwang1009/remove_seq_len

if seq_len node doesn't exist, set sequence_lens of RNN empty

Author: nbcsm

tf2onnx/graph.py

@@ -578,8 +578,8 @@ def update_node_shape_dtype(self, node, override=False):
         # op needs the "Shape" value to infer output shape.
         initializers = []
         for i, inp in enumerate(node.inputs):
-            if not inp:
-                if logger.isEnabledFor(logging.VERBOSE):
+            if inp is None:
+                if logger.isEnabledFor(logging.INFO):
                     logger.warning(
                         "[%s] infer a inexistent node: [%s], please check the code",
                         node.name, node.input[i]

tf2onnx/rewriter/gru_rewriter.py

@@ -182,9 +182,13 @@ def create_rnn_node(self, context):
         context.attributes["direction"] = "forward"
         context.attributes["hidden_size"] = context.hidden_size
         inputs = context.onnx_input_ids
+        # sequence length is optional
+        seq_len_input = utils.ONNX_EMPTY_INPUT
+        if inputs["sequence_lens"]:
+            seq_len_input = inputs["sequence_lens"]
         gru_inputs = [
             inputs["X"], inputs["W"], inputs["R"], inputs["B"],
-            inputs["sequence_lens"], inputs["initial_state"]]
+            seq_len_input, inputs["initial_state"]]
         x_shape = self.g.get_shape(gru_inputs[0])
         x_seq_length = x_shape[0]
         x_batch_size = x_shape[1]

tf2onnx/rewriter/lstm_rewriter.py

@@ -320,9 +320,13 @@ def create_rnn_node(self, context):
         context.attributes["direction"] = "forward"
         context.attributes["hidden_size"] = context.hidden_size
         inputs = context.onnx_input_ids
+        # sequence len input is optional
+        seq_len_input = utils.ONNX_EMPTY_INPUT
+        if inputs["sequence_lens"]:
+            seq_len_input = inputs["sequence_lens"]
         lstm_inputs = [
             inputs["X"], inputs["W"], inputs["R"], inputs["B"],
-            inputs["sequence_lens"], inputs["initial_h"], inputs["initial_c"]]
+            seq_len_input, inputs["initial_h"], inputs["initial_c"]]
 
         x_shape = self.g.get_shape(lstm_inputs[0])
         x_seq_length = x_shape[0]

tf2onnx/rewriter/unit_rnn_rewriter_base.py

@@ -8,9 +8,7 @@
 from __future__ import division
 from __future__ import print_function
 import logging
-from onnx import TensorProto
 
-from tf2onnx.graph_builder import GraphBuilder
 from tf2onnx.rewriter.loop_rewriter_base import LoopRewriterBase, Context
 from tf2onnx.rewriter.rnn_utils import REWRITER_RESULT, get_pattern, \
     get_rnn_scope_name, parse_rnn_loop, seq_len_pattern
@@ -36,7 +34,9 @@ def __init__(self):
         self.hidden_size = None
 
         self.attributes = {} # onnx attributes
-        self.onnx_input_ids = {}  # onnx inputs: [X, W, R, B, sequence_lens, initial_h, initial_c, P]
+        # onnx inputs: [X, W, R, B, sequence_lens, initial_h, initial_c, P],
+        # sequence_lens is optional, i.e., None
+        self.onnx_input_ids = {}
 
 
 class UnitRnnRewriterBase(LoopRewriterBase):
@@ -114,7 +114,9 @@ def parse_unit_rnn(self, context):
         seq_len_node = self.find_sequence_length_node(context)
         if seq_len_node:
             logger.debug("find sequence node: %s", seq_len_node.name)
-            context.seq_len_node = seq_len_node
+            context.onnx_input_ids["sequence_lens"] = seq_len_node.output[0]
+        else:
+            context.onnx_input_ids["sequence_lens"] = None
 
         # require exact one input
         inputs = context.loop_properties.scan_inputs_initial_values
@@ -166,8 +168,6 @@ def rewrite(self, context):
         logger.debug("process the weights/bias/ft_bias, to fit onnx weights/bias requirements")
         self.process_weights_and_bias(context)
 
-        self.process_seq_length(context)
-
         self.process_var_init_nodes(context)
 
         logger.debug("start to build new rnn node")
@@ -222,40 +222,6 @@ def find_sequence_length_node(self, context):
     def process_weights_and_bias(self, context):
         raise NotImplementedError()
 
-    def process_seq_length(self, context):
-        # output: [time step, batch size, input size]
-        seq_len_node = context.seq_len_node
-        shape_node = self.g.make_node("Shape", [context.onnx_input_ids["X"]])
-        # LSTMCell only allow inputs of [batch size, input_size], so we assume dynamic_rnn has 3 dims.
-        # Slice cannot support Int64 in OPSET 7, so we cast here.
-        cast_shape_node = self.g.make_node(
-            "Cast", [shape_node.output[0]],
-            attr={"to": TensorProto.FLOAT},
-            shapes=[self.g.get_shape(shape_node.output[0])]
-        )
-
-        attr = {"axes": [0], "starts": [1], "ends": [2]}
-        inputs_map = {"data": cast_shape_node.output[0], **attr}
-        batchsize_node = GraphBuilder(self.g).make_slice(inputs_map)
-        if not seq_len_node:
-            # Tile's repeats must be INT64
-            repeat_node = self.g.make_node(
-                "Cast", [batchsize_node],
-                attr={"to": TensorProto.INT64}
-            )
-
-            attr = {"axes": [0], "starts": [0], "ends": [1]}
-            inputs_map = {"data": cast_shape_node.output[0], **attr}
-            timestep_node = GraphBuilder(self.g).make_slice(inputs_map)
-            tile_node = self.g.make_node("Tile", [timestep_node, repeat_node.output[0]])
-
-            # LSTM sequence_lens needs to be int32
-            seq_len_node = self.g.make_node(
-                "Cast", [tile_node.output[0]],
-                attr={"to": TensorProto.INT32}
-            )
-        context.onnx_input_ids["sequence_lens"] = seq_len_node.output[0]
-
     def process_var_init_nodes(self, context):
         raise NotImplementedError()
 

tf2onnx/utils.py

@@ -93,6 +93,7 @@ def __init__(self, tensor_id, g):
 
 
 ONNX_UNKNOWN_DIMENSION = -1
+ONNX_EMPTY_INPUT = ""
 
 # index for internally generated names
 INTERNAL_NAME = 1