Why is the loading TF Keras model with LSTM initial states not working?

Question

I am trying the following code to save and load the TF Keras Model which has an LSTM layer with initial_states as inputs. But when I tried to load the model I am getting the error

ValueError: Layer lstm expects 1 inputs, but it received 3 input tensors. Inputs received: [<tf.Tensor 'reshape_1/Identity:0' shape=(None, 5, 5) dtype=float32>, <tf.Tensor 'dense_2/Identity:0' shape=(None, 5) dtype=float32>, <tf.Tensor 'dense_2/Identity:0' shape=(None, 5) dtype=float32>]

Is there any way to load the model with LSTM initial states?

import numpy as np
import tensorflow as tf

inputs = tf.keras.Input(name='input_1', shape=[25])
initial_state = tf.keras.Input(name='initial_state', shape=[5])
dense_x = tf.keras.layers.Dense(units=5)(initial_state)

reshape = tf.keras.layers.Reshape(target_shape=[5, 5])(inputs)

stacked_rnn = tf.keras.layers.LSTM(units=5, return_sequences=True)(inputs=reshape, initial_state=[dense_x, dense_x])
flatten = tf.keras.layers.Flatten()(stacked_rnn)
dense = tf.keras.layers.Dense(name='dense_1', units=1, activation='sigmoid')(flatten)

model = tf.keras.Model(inputs=[inputs, initial_state], outputs=dense, name="test_model")

print(model(inputs=[np.zeros(shape=[5, 25]), np.zeros(shape=[5, 5])]))


tf.keras.models.save_model(model, "sequential/", )

sequential_model = tf.keras.models.load_model("sequential/")

print(sequential_model(inputs=[np.zeros(shape=[5, 25]), np.zeros(shape=[5, 5])]))

TFer TFer · Accepted Answer · 2020-12-30T04:28:02

The above code snippet is working on Tensorflow 2.4.

import numpy as np
import tensorflow as tf

inputs = tf.keras.Input(name='input_1', shape=[25])
initial_state = tf.keras.Input(name='initial_state', shape=[5])
dense_x = tf.keras.layers.Dense(units=5)(initial_state)

reshape = tf.keras.layers.Reshape(target_shape=[5, 5])(inputs)

stacked_rnn = tf.keras.layers.LSTM(units=5, return_sequences=True)(inputs=reshape, initial_state=[dense_x, dense_x])
flatten = tf.keras.layers.Flatten()(stacked_rnn)
dense = tf.keras.layers.Dense(name='dense_1', units=1, activation='sigmoid')(flatten)

model = tf.keras.Model(inputs=[inputs, initial_state], outputs=dense, name="test_model")

#initial states to the model
print(model(inputs=[np.zeros(shape=[5, 25]), np.zeros(shape=[5, 5])]))


tf.keras.models.save_model(model, "sequential/", )

sequential_model = tf.keras.models.load_model("sequential/")
#load model with initial states
print(sequential_model(inputs=[np.zeros(shape=[5, 25]), np.zeros(shape=[5, 5])]))

Output

tf.Tensor(
[[0.5]
 [0.5]
 [0.5]
 [0.5]
 [0.5]], shape=(5, 1), dtype=float32)

tf.Tensor(
[[0.5]
 [0.5]
 [0.5]
 [0.5]
 [0.5]], shape=(5, 1), dtype=float32)

Why is the loading TF Keras model with LSTM initial states not working?

1 Answers