Keep activations in bidirectional LSTM (fixes: onnx#2211)

Signed-off-by: Me <me@example.com>
MaximeChurin · Oct 12, 2023 · 586294b · 586294b
1 parent d5b7f39
commit 586294b
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git a/tests/test_lstm.py b/tests/test_lstm.py
@@ -775,6 +775,31 @@ def func(x):
             return tf.identity(y[0], name="output"), tf.identity(y[1], name="output1")
         self.run_test_case(func, {"input:0": x_val}, [], ["output:0", "output1:0"], rtol=1e-05, atol=1e-06)
 
+    @check_tf_min_version("2.0")
+    @skip_tf_versions("2.1", "Bug in TF 2.1")
+    def test_keras_bilstm_recurrent_activation_is_hard_sigmoid(self):
+        in_shape = [10, 3]
+        x_val = np.random.uniform(size=[2, 10, 3]).astype(np.float32)
+
+        model_in = tf.keras.layers.Input(tuple(in_shape), batch_size=2)
+        x = tf.keras.layers.Bidirectional(
+            tf.keras.layers.LSTM(
+                units=5,
+                return_sequences=True,
+                return_state=True,
+                kernel_initializer=tf.random_uniform_initializer(0.0, 1.0, seed=42),
+                recurrent_initializer=tf.random_uniform_initializer(0.0, 1.0, seed=44),
+                bias_initializer=tf.random_uniform_initializer(0.0, 1.0, seed=43),
+                recurrent_activation="hard_sigmoid",
+            )
+        )(model_in)
+        model = tf.keras.models.Model(inputs=model_in, outputs=x)
+
+        def func(x):
+            y = model(x)
+            return tf.identity(y[0], name="output"), tf.identity(y[1], name="output1")
+        self.run_test_case(func, {"input:0": x_val}, [], ["output:0", "output1:0"], rtol=1e-05, atol=1e-06)
+
     @check_tf_min_version("2.0")
     @skip_tfjs("TFJS converts model incorrectly")
     def test_keras_lstm_sigmoid_dropout(self):

diff --git a/tf2onnx/rewriter/bilstm_rewriter.py b/tf2onnx/rewriter/bilstm_rewriter.py
@@ -54,7 +54,10 @@ def process_bilstm(g, bi_lstms):
         if len(lstm_fw.inputs) > 4:
             lstm_inputs.extend([lstm_fw.input[4], h_node.output[0], c_node.output[0]])
 
-        attr = {"direction": "bidirectional"}
+        attr = {
+            "direction": "bidirectional",
+            "activations": lstm_bw.get_attr_value("activations") + lstm_fw.get_attr_value("activations"),
+        }
         for name in rnn_utils.onnx_rnn_attr_mapping[rnn_utils.ONNX_RNN_TYPE.LSTM]:
             attr_val = lstm_fw.get_attr_value(name)
             if attr_val: