TensorFlow 2.6：无法从保存的模型生成一步预测

Posted 2023-02-16

【中文标题】TensorFlow 2.6：无法从保存的模型生成一步预测

【英文标题】：TensorFlow 2.6: Unable to generate one-step predictions from saved model

【发布时间】：2021-12-09 16:02:06

【问题描述】：

我在尝试使用已保存的 Keras 模型时遇到此错误。

此版本的代码有效。这是一个纸牌游戏的对手手牌预测模型，每个批次大小为 64，每个 25 个时间步长，其中每个步长是一个 211 的暗张量，表示有关前一个时间步长的游戏信息。由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)

这提供了 10 个 shape(1, 156) 张量的输出，正如预期的那样，但是当我保存 HandPredictionModel 时，将其重新加载并使用它来初始化 OneStepModel 我收到有关输入尺寸的错误。

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: 

这可能是什么原因造成的？这里唯一的区别是保存和加载模型的额外步骤。这是一个问题，因为对于我的数据集的大小，我必须以增量方式训练 HandPredictionModel，但任何时候我必须保存和加载它，这意味着我的 OneStepModel 将无法工作。

【参考方案1】：

问题在于，当您保存模型时，用于训练 LSTM 模型的批大小 64 也必须用于预测。可以在此post 中找到此问题的原因以及有关此主题的更多信息。在您的情况下，我通常只是保存模型的权重，然后创建一个新模型并在进行预测时加载权重：

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)

如果您不打算再次训练模型，此选项特别有用。有关此问题的更多替代解决方案，请参阅给定的链接。如果您需要像您提到的那样在迭代中训练您的模型，那么我建议您继续使用save 和load，当您完成训练后，您可以简单地保存权重并加载它们以进行预测。