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

Question

我在尝试使用已保存的 Keras 模型时运行遇到了这个错误。

这个版本的代码有效。它是纸牌游戏的对手手牌预测模型，每个批量大小为 64，每个时间步长为 25，其中每个步骤都是 dim 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)

这给出了 shape(1, 156) 张量的 10 个输出，正如预期的那样，但是当我保存 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 将无法工作。

Answer 1

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

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)

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

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

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

python

tensorflow

keras

tensorflow2.0

tf.keras