diff --git a/keras2onnx/ke2onnx/bidirectional.py b/keras2onnx/ke2onnx/bidirectional.py index f4cbc934..bc2eee2d 100644 --- a/keras2onnx/ke2onnx/bidirectional.py +++ b/keras2onnx/ke2onnx/bidirectional.py @@ -240,10 +240,17 @@ def convert_bidirectional(scope, operator, container): apply_split(scope, transposed_y_name, [forward_y_name, backward_y_name], container, axis=2) # Change (T, N, 1, C') into (T, N, C') to meet Keras spec - container.add_node('Squeeze', forward_y_name, operator.outputs[0].full_name, + forward_y_name_1 = scope.get_unique_variable_name(operator.full_name + '_Y_forward_1') + backward_y_name_1 = scope.get_unique_variable_name(operator.full_name + '_Y_backward_1') + container.add_node('Squeeze', forward_y_name, forward_y_name_1, name=scope.get_unique_variable_name('Squeeze'), axes=[2]) - container.add_node('Squeeze', backward_y_name, operator.outputs[1].full_name, + container.add_node('Squeeze', backward_y_name, backward_y_name_1, name=scope.get_unique_variable_name('Squeeze'), axes=[2]) + + apply_reshape(scope, forward_y_name_1, operator.outputs[0].full_name, container, + desired_shape=[-1, seq_length, hidden_size]) + apply_reshape(scope, backward_y_name_1, operator.outputs[1].full_name, container, + desired_shape=[-1, seq_length, hidden_size]) else: perm = [1, 0, 2] if merge_concat: diff --git a/tests/test_layers.py b/tests/test_layers.py index 45df6324..ddd5a052 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -642,19 +642,34 @@ def test_LSTM_reshape(self): self.assertTrue(self.run_onnx_runtime('tf_lstm', onnx_model, data, expected)) def test_Bidirectional(self): - input_dim = 10 - sequence_len = 5 - model = keras.Sequential() - model.add(keras.layers.Bidirectional(keras.layers.LSTM(10, return_sequences=False), - input_shape=(5, 10))) - model.add(keras.layers.Dense(5)) - model.add(keras.layers.Activation('softmax')) - model.compile(loss='categorical_crossentropy', optimizer='rmsprop') - - onnx_model = keras2onnx.convert_keras(model, 'test') - data = np.random.rand(input_dim, sequence_len).astype(np.float32).reshape((1, sequence_len, input_dim)) - expected = model.predict(data) - self.assertTrue(self.run_onnx_runtime('bidirectional', onnx_model, data, expected)) + for return_sequences in [True, False]: + input_dim = 10 + sequence_len = 5 + model = keras.Sequential() + model.add(keras.layers.Bidirectional(keras.layers.LSTM(10, return_sequences=return_sequences), + input_shape=(5, 10))) + model.add(keras.layers.Dense(5)) + model.add(keras.layers.Activation('softmax')) + model.compile(loss='categorical_crossentropy', optimizer='rmsprop') + + onnx_model = keras2onnx.convert_keras(model, 'test') + data = np.random.rand(input_dim, sequence_len).astype(np.float32).reshape((1, sequence_len, input_dim)) + expected = model.predict(data) + self.assertTrue(self.run_onnx_runtime('bidirectional', onnx_model, data, expected)) + + for merge_mode in ['concat', None]: + # TODO: case return_sequences=False + for return_sequences in [True]: + input_dim = 10 + sequence_len = 5 + sub_input1 = keras.layers.Input(shape=(sequence_len, input_dim)) + sub_mapped1 = keras.layers.Bidirectional(keras.layers.LSTM(10, return_sequences=return_sequences), + input_shape=(5, 10), merge_mode=merge_mode)(sub_input1) + keras_model = keras.Model(inputs=sub_input1, outputs=sub_mapped1) + onnx_model = keras2onnx.convert_keras(keras_model, 'test_2') + data = np.random.rand(input_dim, sequence_len).astype(np.float32).reshape((1, sequence_len, input_dim)) + expected = keras_model.predict(data) + self.assertTrue(self.run_onnx_runtime('bidirectional', onnx_model, data, expected)) def test_separable_convolution(self): N, C, H, W = 2, 3, 5, 5