TensorFlow 2.0: Save and load a model that contains a LSTM layer, while the load commond failed with ValueError

In case anyone else stumbles upon this, this solution worked for me:

# Save model
tf.keras.models.save_model(model, "saved_model.hp5", save_format="h5")

# Load model
loaded_model = tf.keras.models.load_model("saved_model.hp5")

Not sure why the "model.save(filename)" syntax doesn't work with LSTM, but I ran into the same issue.

TL;DR try providing a default value for training, i.e.,

def call(self, inputs, training=None):

I encountered similar error with tensorflow 2.1.0:

"ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (2 total):
    * Tensor("inputs:0", shape=(None, 128, 128, 3), dtype=float32)
    * Tensor("training:0", shape=(), dtype=bool)

My model does not have LSTM layers. From the fact that the complaint in the error message is related to inputs and training, I guess it's related to the call function that is required when we subclass the Model class, since inputs and training are two of the argument names for that function. What solved my problem is:

instead of

def call(self, inputs):

provide the default value for training, i.e.,

def call(self, inputs, training=None):

There is a difference between subclass model and other keras model types(Sequential and Functional). Due to those differences listed here, saving and loading of subclass models is different from other Keras models.

Important points

1. A subclassed model differs in that it's not a datastructure, it's a piece of code. The architecture of the model is defined via the body of the call method. This means that the architecture of the model cannot be safely serialized.
2. A subclassed model that has never been used cannot be saved. That's because a subclassed model needs to be called on some data in order to create its weights. Until the model has been called, it does not know the shape and dtype of the input data it should be expecting, and thus cannot create its weight variables.

Having mentioned about differences between subclass and other keras models (Sequential and functional), there are three different ways one can save a subcalssed model.

let's create a model

class ThreeLayerMLP(keras.Model):

  def __init__(self, name=None):
    super(ThreeLayerMLP, self).__init__(name=name)
    self.dense_1 = layers.Dense(64, activation='relu', name='dense_1')
    self.dense_2 = layers.Dense(64, activation='relu', name='dense_2')
    self.pred_layer = layers.Dense(10, name='predictions')

  def call(self, inputs):
    x = self.dense_1(inputs)
    x = self.dense_2(x)
    return self.pred_layer(x)

def get_model():
  return ThreeLayerMLP(name='3_layer_mlp')

model = get_model()

train the model

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255
x_test = x_test.reshape(10000, 784).astype('float32') / 255

model.compile(loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              optimizer=keras.optimizers.RMSprop())
history = model.fit(x_train, y_train,
                    batch_size=64,
                    epochs=1)
# Reset metrics before saving so that loaded model has same state,
# since metric states are not preserved by Model.save_weights
model.reset_metrics()

Approach 1: Simplest Approach

This approach is by using model.save to save whole model and by using load_model to restore previously stored subclassed model.

# Save the model
model.save('path_to_my_model',save_format='tf')

# Recreate the exact same model purely from the file
new_model = keras.models.load_model('path_to_my_model')

Approach 2: Identical to Approach 1

Second approach is by using tf.saved_model.save. This is equivalent to the tf format in model.save. You can once again call load_model to restore the previously saved subclassed model.

Approach 3:

Third approach is to use save_weights to create a TensorFlow SavedModel checkpoint. To restore a subclass model, (1) you will need access to the code that created the model object and create the model again, (2) compile the model in order to restore the optimizer state and any stateful metrics, and (3) call it on some data before calling load_weights.

model.save_weights('path_to_my_weights', save_format='tf')

# Recreate the model
new_model = get_model()
new_model.compile(loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  optimizer=keras.optimizers.RMSprop())

# This initializes the variables used by the optimizers,
# as well as any stateful metric variables
new_model.train_on_batch(x_train[:1], y_train[:1])

# Load the state of the old model
new_model.load_weights('path_to_my_weights')

Another important things to note while saving the model is as follows.

Entire model can be saved in two different file formats (SavedModel and HDF5). It is to be noted that TensorFlow SavedModel ('tf') format is the default file format in TF2.x. However, model can be saved in HDF5 format. The key difference between HDF5 and SavedModel is that HDF5 uses object configs to save the model architecture, while SavedModel saves the execution graph. Thus, SavedModels (models that are saved with save_format='tf') are able to save custom objects like subclassed models and custom layers without requiring the original code.

Hope this helps.