Skip to content

Commit ca17b3c

Browse files
Keep activations in bidirectional LSTM (fixes: #2211) (#2234)
* Keep activations in bidirectional LSTM (fixes: #2211) --------- Signed-off-by: Me <[email protected]> Co-authored-by: Jay Zhang <[email protected]>
1 parent b57fe91 commit ca17b3c

File tree

2 files changed

+33
-1
lines changed

2 files changed

+33
-1
lines changed

tests/test_lstm.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -775,6 +775,30 @@ def func(x):
775775
return tf.identity(y[0], name="output"), tf.identity(y[1], name="output1")
776776
self.run_test_case(func, {"input:0": x_val}, [], ["output:0", "output1:0"], rtol=1e-05, atol=1e-06)
777777

778+
@check_tf_min_version("2.0")
779+
def test_keras_bilstm_recurrent_activation_is_hard_sigmoid(self):
780+
in_shape = [10, 3]
781+
x_val = np.random.uniform(size=[2, 10, 3]).astype(np.float32)
782+
783+
model_in = tf.keras.layers.Input(tuple(in_shape), batch_size=2)
784+
x = tf.keras.layers.Bidirectional(
785+
tf.keras.layers.LSTM(
786+
units=5,
787+
return_sequences=True,
788+
return_state=True,
789+
kernel_initializer=tf.random_uniform_initializer(0.0, 1.0, seed=42),
790+
recurrent_initializer=tf.random_uniform_initializer(0.0, 1.0, seed=44),
791+
bias_initializer=tf.random_uniform_initializer(0.0, 1.0, seed=43),
792+
recurrent_activation="hard_sigmoid",
793+
)
794+
)(model_in)
795+
model = tf.keras.models.Model(inputs=model_in, outputs=x)
796+
797+
def func(x):
798+
y = model(x)
799+
return tf.identity(y[0], name="output"), tf.identity(y[1], name="output1")
800+
self.run_test_case(func, {"input:0": x_val}, [], ["output:0", "output1:0"], rtol=1e-05, atol=1e-06)
801+
778802
@check_tf_min_version("2.0")
779803
@skip_tfjs("TFJS converts model incorrectly")
780804
def test_keras_lstm_sigmoid_dropout(self):

tf2onnx/rewriter/bilstm_rewriter.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,11 +54,19 @@ def process_bilstm(g, bi_lstms):
5454
if len(lstm_fw.inputs) > 4:
5555
lstm_inputs.extend([lstm_fw.input[4], h_node.output[0], c_node.output[0]])
5656

57-
attr = {"direction": "bidirectional"}
57+
direction = "bidirectional"
58+
attr = {"direction": direction}
5859
for name in rnn_utils.onnx_rnn_attr_mapping[rnn_utils.ONNX_RNN_TYPE.LSTM]:
5960
attr_val = lstm_fw.get_attr_value(name)
6061
if attr_val:
6162
attr[name] = attr_val
63+
# activation has to be took care, attr here is proto
64+
activations = [act.decode("utf-8")
65+
for act in lstm_fw.get_attr_value("activations", [])]
66+
activations += [act.decode("utf-8")
67+
for act in lstm_bw.get_attr_value("activations", [])]
68+
if activations:
69+
attr["activations"] = activations
6270

6371
bi_lstm_node = g.make_node("LSTM", lstm_inputs, attr=attr, output_count=3)
6472
all_nodes.append(bi_lstm_node)

0 commit comments

Comments
 (0)