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

user_3081907 0
Viewed 6516

When I was trying to save and load a model that contains a LSTM layer, the load commond failed with ValueError: Could not find matching function to call loaded from the SavedModel.

class RegNet(Model):
    def __init__(self,
             intermediate_dim=50,
             state_dim=9,
             name='RegNet',
             **kwargs):
        super(RegNet, self).__init__()
        self.d1 = Dense(intermediate_dim, activation='relu')
        self.d2 = Dense(state_dim, activation='relu')
        self.h = LSTM(state_dim, activation='sigmoid', return_sequences=True)
        self.o = Dense(state_dim, activation='softmax')

    def call(self, x):
        x = self.d1(x)
        x = self.d2(x)
        x = self.h(x)
        y = self.o(x)
        return y

regNet = RegNet()
...
# Export the model to a SavedModel
regNet.save(regNet_ckpt_dir, save_format='tf')
# Recreate the exact same model
tf.keras.models.load_model(regNet_ckpt_dir)

The error report:

> ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (2 total):
    * Tensor("x:0", shape=(None, 1, 20), dtype=float32)
    * Tensor("training:0", shape=(), dtype=bool)
  Keyword arguments: {}

Expected these arguments to match one of the following 4 option(s):

Option 1:
  Positional arguments (2 total):
    * TensorSpec(shape=(None, 1, 20), dtype=tf.float32, name='input_1')
    * False
  Keyword arguments: {}

Option 2:
  Positional arguments (2 total):
    * TensorSpec(shape=(None, 1, 20), dtype=tf.float32, name='x')
    * False
  Keyword arguments: {}

Option 3:
  Positional arguments (2 total):
    * TensorSpec(shape=(None, 1, 20), dtype=tf.float32, name='x')
    * True
  Keyword arguments: {}

Option 4:
  Positional arguments (2 total):
    * TensorSpec(shape=(None, 1, 20), dtype=tf.float32, name='input_1')
    * True
  Keyword arguments: {}

When I comment the LSTM layer, the load command will succeed. Where is the problem? We cannot save and load a model with LSTM layer in TensorFlow 2.0?

tensorflow2.0
4 Answers

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.

user_122353250

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):
user_73823140

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.

user_99362280

Looks like, adding a compute_output_shape resolves the issue.

regNet.compute_output_shape(input_shape=(1,10,3))

Here is a gist for reference.

I have used TF2.7 and the model was successfully saved and loaded.

user_142906810
Related