TensorFlow 2.6: Unable to generate one-step predictions from saved model

Question

I am running into this error when trying to use a saved Keras model.

This version of the code works. It is an opponent hand prediction model for a card game that takes in a batch size 64 of 25 timesteps each, where each step is a tensor of dim 211 that represents information about the game in the previous timestep. It's modified from this official tutorial.

# Main model
class HandPredictionModel(tf.keras.Model):

    def __init__(self):
        super().__init__(self)
        self.lstm1 = tf.keras.layers.LSTM(512, return_sequences=True)
        self.dropout1 = tf.keras.layers.Dropout(0.2)
        self.lstm2 = tf.keras.layers.LSTM(512, return_sequences=True, return_state=True)
        self.dense = tf.keras.layers.Dense(156, activation="sigmoid")
            
    @tf.function
    def call(self, x, states=None, return_state=False, training=False):
        if states is None:
            states = self.lstm1.get_initial_state(x)
        x = self.lstm1(x, states)
        x = self.dropout1(x)
        x, final_memory_state, final_carry_state = self.lstm2(x)
        x = self.dense(x)
        if return_state:
            return x, final_memory_state, final_carry_state
        return x

handPredictionmodel = HandPredictionModel()
handPredictionModel.compile(...) # loss function, optimizer

dataset = (dataset.shuffle(1000, reshuffle_each_iteration=True).batch(64, drop_remainder=True)) 
# <BatchDataset shapes: ((64, 25, 211), (64, 25, 156)), types: (tf.float32, tf.float32)>

history = handPredictionModel.fit(dataset, epochs=100)

# One-step model
class OneStep(tf.keras.Model):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    @tf.function
    def predict(self, inputs, states=None):
        inputs = tf.expand_dims(tf.expand_dims(inputs, axis=0), axis=0) # add 'fake' dims for batch and timestep
        predicted_logits, memory_state, carry_state= self.model(x=inputs, states=states, return_state=True, training=False)
        predicted_logits = predicted_logits[:, -1, :]

        return predicted_logits, [memory_state, carry_state]

# Testing
oneStepModel = OneStep(handPredictionModel)

states = None
for i in range(10):
    t = tf.zeros([211])
    pred, states = oneStepModel.predict(t, states)
    print(pred)

This gives 10 outputs of shape(1, 156) tensors, as expected, however when I save HandPredictionModel, load it back in, and use that to initialize OneStepModel I get this error about input dimensions.

tf.saved_model.save(model, 'handPredictionModel')
loadedModel = tf.saved_model.load('handPredictionModel')
oneStepModel = OneStep(loadedModel)

states = None
for i in range(10):
    t = tf.zeros([211])
    pred, states = oneStepModel.predict(t, states)
    print(pred)

    ValueError: Could not find matching function to call loaded from the SavedModel. Got:
      Positional arguments (4 total):
        * Tensor("x:0", shape=(1, 1, 211), dtype=float32)
        * None
        * True
        * False
      Keyword arguments: {}

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

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

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

    Option 3:
      Positional arguments (4 total):
        * TensorSpec(shape=(None, 25, 211), dtype=tf.float32, name='x')
        * None
        * False
        * True
      Keyword arguments: {}

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

What could be causing this? The only difference here is the extra step of saving and loading the model. This is a problem because with the size of my dataset I have to train HandPredictionModel in increments, but any time I have to save and load it this means my OneStepModel will not work.

AloneTogether · Accepted Answer · 2021-10-23T10:02:52.087

The problem is that when you save your model, the batch size 64 used to train your LSTM model must also be used for prediction. The reason for this and more information on this topic can be found in this post. In your case, I usually just save the weights of the model, then create a new model and load the weights when making predictions:

class OneStep(tf.keras.Model):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    @tf.function
    def predict(self, inputs, states=None):
        inputs = tf.expand_dims(tf.expand_dims(inputs, axis=0), axis=0) # add 'fake' dims for batch and timestep
        predicted_logits, memory_state, carry_state= self.model.call(x=inputs, states=states, return_state=True, training=False)
        predicted_logits = predicted_logits[:, -1, :]

        return predicted_logits, [memory_state, carry_state]

handPredictionModel.save_weights('model_weights') 
loadedModel = HandPredictionModel()
loadedModel.load_weights('model_weights')
oneStepModel = OneStep(loadedModel)

states = None
for i in range(10):
    t = tf.zeros([211])
    pred, states = oneStepModel.predict(t, states)
    print(pred)

This option is especially useful if you do not plan to train your model again. See the given link for more alternative solutions to this problem. If you need to train your model in iterations, as you mentioned, then I would recommend that you continue to use save and load and when you are done training, you can simply save the weights and load them for predictions.

TensorFlow 2.6: Unable to generate one-step predictions from saved model

1 Answers1