Merge pull request #961 from chinhuang007/support-gru-14

Add batchwise support for gru, lstm, rnn

Author: chudegao

doc/support_status.md

@@ -1,7 +1,7 @@
 # ONNX-Tensorflow Support Status
 |||
 |-:|:-|
-|ONNX-Tensorflow Version|Master ( commit id: 9ab9b934c2c8494b6309d20f15acabcb3abd126d )|
+|ONNX-Tensorflow Version|Master ( commit id: a1005fbd2a95699a34f83d5d25fe20d5213860d3 )|
 |ONNX Version|Master ( commit id: 1f63dcb7fcc3a8bf5c3c8e326867ecd6f5c43f35 )|
 |Tensorflow Version|v2.5.0|
 
@@ -62,7 +62,7 @@ Notes:
 |EyeLike|-|-|-|-|-|-|-|-|**9**|9|9|9|9|9|9|EyeLike|
 |Flatten|**1**|1|1|1|1|1|1|1|**9**|9|**11**|11|**13**|13|13|Flatten|
 |Floor|**1**|1|1|1|1|**6**|6|6|6|6|6|6|**13**|13|13|Floor|
-|GRU|**1**:small_orange_diamond:|1:small_orange_diamond:|**3**:small_orange_diamond:|3:small_orange_diamond:|3:small_orange_diamond:|3:small_orange_diamond:|**7**:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|**14**:small_red_triangle:|14:small_red_triangle:|GRU|
+|GRU|**1**:small_orange_diamond:|1:small_orange_diamond:|**3**:small_orange_diamond:|3:small_orange_diamond:|3:small_orange_diamond:|3:small_orange_diamond:|**7**:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|**14**:small_orange_diamond:|14:small_orange_diamond:|GRU|
 |Gather|**1**|1|1|1|1|1|1|1|1|1|**11**|11|**13**|13|13|Gather|
 |GatherElements|-|-|-|-|-|-|-|-|-|-|**11**|11|**13**|13|13|GatherElements|
 |GatherND|-|-|-|-|-|-|-|-|-|-|**11**|**12**|**13**|13|13|GatherND|
@@ -81,7 +81,7 @@ Notes:
 |IsInf|-|-|-|-|-|-|-|-|-|**10**|10|10|10|10|10|IsInf|
 |IsNaN|-|-|-|-|-|-|-|-|**9**|9|9|9|**13**|13|13|IsNaN|
 |LRN|**1**|1|1|1|1|1|1|1|1|1|1|1|**13**|13|13|LRN|
-|LSTM|**1**:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|**7**:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|**14**:small_red_triangle:|14:small_red_triangle:|LSTM|
+|LSTM|**1**:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|**7**:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|**14**:small_orange_diamond:|14:small_orange_diamond:|LSTM|
 |LeakyRelu|**1**|1|1|1|1|**6**|6|6|6|6|6|6|6|6|6|LeakyRelu|
 |Less|**1**|1|1|1|1|1|**7**|7|**9**|9|9|9|**13**|13|13|Less|
 |LessOrEqual|-|-|-|-|-|-|-|-|-|-|-|**12**|12|12|12|LessOrEqual|
@@ -118,7 +118,7 @@ Notes:
 |QLinearConv|-|-|-|-|-|-|-|-|-|**10**|10|10|10|10|10|QLinearConv|
 |QLinearMatMul|-|-|-|-|-|-|-|-|-|**10**|10|10|10|10|10|QLinearMatMul|
 |QuantizeLinear|-|-|-|-|-|-|-|-|-|**10**|10|10|**13**|13|13|QuantizeLinear|
-|RNN|**1**:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|**7**:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|**14**:small_red_triangle:|14:small_red_triangle:|RNN|
+|RNN|**1**:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|1:small_orange_diamond:|**7**:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|7:small_orange_diamond:|**14**:small_orange_diamond:|14:small_orange_diamond:|RNN|
 |RandomNormal|**1**|1|1|1|1|1|1|1|1|1|1|1|1|1|1|RandomNormal|
 |RandomNormalLike|**1**|1|1|1|1|1|1|1|1|1|1|1|1|1|1|RandomNormalLike|
 |RandomUniform|**1**|1|1|1|1|1|1|1|1|1|1|1|1|1|1|RandomUniform|
@@ -186,7 +186,7 @@ Notes:
 |Where|-|-|-|-|-|-|-|-|**9**|9|9|9|9|9|9|Where|
 |Xor|**1**|1|1|1|1|1|**7**|7|7|7|7|7|7|7|7|Xor|
 
-ONNX-TF Supported Operators / ONNX Operators: 151 / 169
+ONNX-TF Supported Operators / ONNX Operators: 154 / 169
 
 Notes:
 1. BatchNormalization: BatchNormalization with training_mode=1 is not supported in Tensorflow converte.

onnx_tf/handlers/backend/gru.py

@@ -161,6 +161,11 @@ def _common(cls, node, **kwargs):
     hidden_size = node.attrs["hidden_size"]
     direction = node.attrs.get("direction", "forward")
     num_directions = 2 if direction == "bidirectional" else 1
+    layout = node.attrs.get("layout", 0)
+
+    # Need transpose for batchwise
+    if layout == 1:
+      x = tf.transpose(x, perm=[1, 0, 2])
 
     # removed from version 7, default is 0
     output_sequence = node.attrs.get("output_sequence", 0)
@@ -207,6 +212,9 @@ def _common(cls, node, **kwargs):
       if input_size == 6:
         initial_h = tensor_dict.get(node.inputs[5], None)
         if initial_h is not None:
+          # Need transpose for batchwise
+          if layout == 1:
+            initial_h = tf.transpose(initial_h, perm=[1, 0, 2])
           initial_state = (initial_h[0],)
           if num_directions == 2:
             initial_state_bw = (initial_h[1],)
@@ -241,6 +249,11 @@ def _common(cls, node, **kwargs):
       output_bw = tf.expand_dims(output_bw, 1)
       output = tf.concat((output_fw, output_bw), axis=1)
 
+    # Need transpose for batchwise
+    if layout == 1:
+      output = tf.transpose(output, perm=[2, 0, 1, 3])
+      h = tf.transpose(h, perm=[1, 0, 2])
+
     return [output, h] if output_sequence == 0 else [h]
 
   @classmethod
@@ -254,3 +267,7 @@ def version_3(cls, node, **kwargs):
   @classmethod
   def version_7(cls, node, **kwargs):
     return cls._common(node, **kwargs)
+
+  @classmethod
+  def version_14(cls, node, **kwargs):
+    return cls._common(node, **kwargs)

onnx_tf/handlers/backend/lstm.py

@@ -168,6 +168,11 @@ def _common(cls, node, **kwargs):
     hidden_size = node.attrs["hidden_size"]
     direction = node.attrs.get("direction", "forward")
     num_directions = 2 if direction == "bidirectional" else 1
+    layout = node.attrs.get("layout", 0)
+
+    # Need transpose for batchwise
+    if layout == 1:
+      x = tf.transpose(x, perm=[1, 0, 2])
 
     # removed from version 7, default is 0
     output_sequence = node.attrs.get("output_sequence", 0)
@@ -220,6 +225,10 @@ def _common(cls, node, **kwargs):
         initial_c = tensor_dict.get(
             node.inputs[6],
             None) if input_size >= 7 else tf.zeros_like(initial_h)
+        # Need transpose for batchwise
+        if layout == 1:
+            initial_h = tf.transpose(initial_h, perm=[1, 0, 2])
+            initial_c = tf.transpose(initial_c, perm=[1, 0, 2])
         if initial_h is not None and initial_c is not None:
           initial_state = (tf.compat.v1.nn.rnn_cell.LSTMStateTuple(
               initial_c[0], initial_h[0]),)
@@ -261,6 +270,12 @@ def _common(cls, node, **kwargs):
       output_bw = tf.expand_dims(output_bw, 1)
       output = tf.concat((output_fw, output_bw), axis=1)
 
+    # Need transpose for batchwise
+    if layout == 1:
+      output = tf.transpose(output, perm=[2, 0, 1, 3])
+      h = tf.transpose(h, perm=[1, 0, 2])
+      c = tf.transpose(c, perm=[1, 0, 2])
+
     return [output, h, c] if output_sequence == 0 else [h, c]
 
   @classmethod
@@ -270,3 +285,7 @@ def version_1(cls, node, **kwargs):
   @classmethod
   def version_7(cls, node, **kwargs):
     return cls._common(node, **kwargs)
+
+  @classmethod
+  def version_14(cls, node, **kwargs):
+    return cls._common(node, **kwargs)

onnx_tf/handlers/backend/rnn.py

@@ -104,6 +104,11 @@ def _common(cls, node, **kwargs):
     direction = node.attrs.get("direction", "forward")
     num_directions = 2 if direction == "bidirectional" else 1
     output_sequence = node.attrs.get("output_sequence", 0)
+    layout = node.attrs.get("layout", 0)
+
+    # Need transpose for batchwise
+    if layout == 1:
+      x = tf.transpose(x, perm=[1, 0, 2])
 
     # TODO(fumihwh): check if prev node is one of RNN
     # process input if it comes from other previous cell
@@ -145,6 +150,9 @@ def _common(cls, node, **kwargs):
       if input_size == 6:
         initial_h = tensor_dict.get(node.inputs[5], None)
         if initial_h is not None:
+          # Need transpose for batchwise
+          if layout == 1:
+            initial_h = tf.transpose(initial_h, perm=[1, 0, 2])
           initial_state = (initial_h[0],)
           if num_directions == 2:
             initial_state_bw = (initial_h[1],)
@@ -179,6 +187,11 @@ def _common(cls, node, **kwargs):
       output_bw = tf.expand_dims(output_bw, 1)
       output = tf.concat((output_fw, output_bw), axis=1)
 
+    # Need transpose for batchwise
+    if layout == 1:
+      output = tf.transpose(output, perm=[2, 0, 1, 3])
+      h = tf.transpose(h, perm=[1, 0, 2])
+
     return [output, h] if output_sequence == 0 else [h]
 
   @classmethod
@@ -188,3 +201,7 @@ def version_1(cls, node, **kwargs):
   @classmethod
   def version_7(cls, node, **kwargs):
     return cls._common(node, **kwargs)
+
+  @classmethod
+  def version_14(cls, node, **kwargs):
+    return cls._common(node, **kwargs)

onnx_tf/opset_version.py

@@ -54,7 +54,7 @@
     'FeatureVectorizer': [],
     'Flatten': [1, 9, 11, 13],
     'Floor': [1, 6, 13],
-    'GRU': [1, 3, 7],
+    'GRU': [1, 3, 7, 14],
     'Gather': [1, 11, 13],
     'GatherElements': [11, 13],
     'GatherND': [11, 12, 13],
@@ -76,7 +76,7 @@
     'IsInf': [10],
     'IsNaN': [9, 13],
     'LRN': [1, 13],
-    'LSTM': [1, 7],
+    'LSTM': [1, 7, 14],
     'LabelEncoder': [],
     'LeakyRelu': [1, 6],
     'Less': [1, 7, 9, 13],
@@ -119,7 +119,7 @@
     'QLinearConv': [10],
     'QLinearMatMul': [10],
     'QuantizeLinear': [10, 13],
-    'RNN': [1, 7],
+    'RNN': [1, 7, 14],
     'RandomNormal': [1],
     'RandomNormalLike': [1],
     'RandomUniform': [1],