TensorFlow 2.6 加载已保存Keras模型后无法生成单步预测问题

• 你为call方法添加的@tf.function装饰器只会追踪实际触发过的参数组合对应的计算图，训练阶段你只给模型喂过形状为(None, 25, 211)的输入，且从未设置过return_state=True调用call方法，因此保存到SavedModel中的函数签名只有报错信息里列出的4种，不支持推理时(1, 1, 211)输入、return_state=True的调用场景。
• tf.saved_model.save保存的是静态计算图签名，加载后得到的是封装好的调用对象，不会保留原自定义Keras类的动态调度逻辑，只能严格匹配已保存的签名调用。

方案一（最适合分阶段训练场景） #

改用Keras权重保存加载方案，保留原类的所有动态逻辑：

# 训练完成后保存权重
handPredictionModel.save_weights('hand_prediction_weights.h5')

# 需要加载时，先实例化原模型类，再加载权重
loadedModel = HandPredictionModel()
loadedModel.compile(...) # 和训练时的编译配置一致即可继续训练
loadedModel.load_weights('hand_prediction_weights.h5')

# 后续初始化OneStepModel的逻辑和原来完全一致
oneStepModel = OneStep(loadedModel)

方案二（必须使用SavedModel格式时） #

提前为call方法指定全场景签名，并在保存前触发推理场景的调用，确保对应计算图被追踪：

1. 修改HandPredictionModel的call方法装饰器，指定动态输入签名：

class HandPredictionModel(tf.keras.Model):
    # __init__逻辑不变
    @tf.function(input_signature=[
        tf.TensorSpec(shape=(None, None, 211), dtype=tf.float32, name='x'),
        tf.Optional([tf.TensorSpec(shape=(None, 512), dtype=tf.float32), tf.TensorSpec(shape=(None, 512), dtype=tf.float32)]),
        tf.TensorSpec(shape=(), dtype=tf.bool, name='return_state'),
        tf.TensorSpec(shape=(), dtype=tf.bool, name='training')
    ])
    def call(self, x, states=None, return_state=False, training=False):
        # 原逻辑不变

2. 保存模型前手动触发一次推理场景的调用，完成对应计算图的追踪：

# 训练完成后，先触发一次推理场景调用
_ = handPredictionModel(tf.random.normal((1, 1, 211)), return_state=True)
# 再保存模型
tf.saved_model.save(handPredictionModel, 'handPredictionModel')

修改后加载的模型即可正常支持推理场景的调用。

内容的提问来源于stack exchange，提问作者lumos309