加载TensorFlow循环神经网络模型报错：调用my_model层时遇异常

使用TensorFlow与Python基于循环神经网络（RNN）生成语录，模型训练及保存过程无报错，但在另一个脚本中加载模型时触发ValueError，提示无法找到匹配的SavedModel具体调用函数。模型保存目录包含子文件夹与.pb文件，附错误信息及加载代码。

错误信息 #

ValueError: Exception encountered when calling layer "my_model" "                 f"(type MyModel).

Could not find matching concrete function to call loaded from the SavedModel. Got:
  Positional arguments (4 total):
    * <tf.Tensor 'inputs:0' shape=(1, None) dtype=int64>
    * 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, 100), dtype=tf.int64, name='inputs')
    * None
    * False
    * False
  Keyword arguments: {}

Option 2:
  Positional arguments (4 total):
    * TensorSpec(shape=(None, 100), dtype=tf.int64, name='inputs')
    * None
    * False
    * True
  Keyword arguments: {}

Option 3:
  Positional arguments (4 total):
    * TensorSpec(shape=(None, 100), dtype=tf.int64, name='input_1')
    * None
    * False
    * False
  Keyword arguments: {}

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

Call arguments received by layer "my_model" "                 f"(type MyModel):
  • args=('tf.Tensor(shape=(1, None), dtype=int64)',)
  • kwargs={'states': 'None', 'return_state': 'True', 'training': 'False'}

加载模型代码 #

import tensorflow as tf
import time

path_to_file = tf.keras.utils.get_file('data.txt', 'file:///C:/Users/sebas/Desktop/quites/data.txt')
text = open(path_to_file, 'rb').read().decode(encoding='utf-8')
vocab = sorted(set(text))
ids_from_chars = tf.keras.layers.StringLookup(
    vocabulary=list(vocab), mask_token=None)
chars_from_ids = tf.keras.layers.StringLookup(
    vocabulary=ids_from_chars.get_vocabulary(), invert=True, mask_token=None)

new_model = tf.keras.models.load_model('quoteModel')
class OneStep(tf.keras.Model):
  def __init__(self, model, chars_from_ids, ids_from_chars, temperature=1.0):
    super().__init__()
    self.temperature = temperature
    self.model = model
    self.chars_from_ids = chars_from_ids
    self.ids_from_chars = ids_from_chars
    skip_ids = self.ids_from_chars(['[UNK]'])[:, None]
    sparse_mask = tf.SparseTensor(
        values=[-float('inf')]*len(skip_ids),
        indices=skip_ids,
        dense_shape=[len(ids_from_chars.get_vocabulary())])
    self.prediction_mask = tf.sparse.to_dense(sparse_mask)

  @tf.function
  def generate_one_step(self, inputs, states=None):
    input_chars = tf.strings.unicode_split(inputs, 'UTF-8')
    input_ids = self.ids_from_chars(input_chars).to_tensor()
    predicted_logits, states = self.model(inputs=input_ids, states=states,
                                          return_state=True)
    predicted_logits = predicted_logits[:, -1, :]
    predicted_logits = predicted_logits/self.temperature
    predicted_logits = predicted_logits + self.prediction_mask
    predicted_ids = tf.random.categorical(predicted_logits, num_samples=1)
    predicted_ids = tf.squeeze(predicted_ids, axis=-1)
    predicted_chars = self.chars_from_ids(predicted_ids)
    return predicted_chars, states

one_step_model = OneStep(new_model, chars_from_ids, ids_from_chars)
start = time.time()
states = None
next_char = tf.constant(['I'])
result = [next_char]

for n in range(1000):
  next_char, states = one_step_model.generate_one_step(next_char, states=states)
  result.append(next_char)

result = tf.strings.join(result)
end = time.time()
print(result[0].numpy().decode('utf-8'), '\n\n' + '_'*80)

print(one_step_model.generate_one_step("I"))

解决方案 #

1. 核心问题分析 #

错误显示模型训练时接受的输入形状是(None, 100)（固定序列长度100），但生成时传入的输入形状是(1, None)（变长序列），同时调用时指定的return_state=True在模型保存的签名中匹配逻辑异常，导致找不到可用的调用函数。

2. 具体修复方法 #

方法一：调整生成代码的输入形状

在generate_one_step方法中，将输入序列强制调整为训练时的固定长度100：

@tf.function
def generate_one_step(self, inputs, states=None):
    input_chars = tf.strings.unicode_split(inputs, 'UTF-8')
    input_ids = self.ids_from_chars(input_chars).to_tensor()
    # 填充或截断到固定长度100
    seq_len = tf.shape(input_ids)[1]
    if seq_len < 100:
        input_ids = tf.pad(input_ids, [[0, 0], [0, 100 - seq_len]])
    input_ids = tf.slice(input_ids, [0, 0], [-1, 100])
    # 后续逻辑不变
    predicted_logits, states = self.model(inputs=input_ids, states=states,
                                          return_state=True)
    ...

方法二：确保模型保存时包含生成所需签名

训练完成后，保存模型前先触发一次带return_state=True的调用，让TensorFlow记录该签名：

# 训练完成后执行
test_input = tf.random.uniform((1, 100), maxval=len(vocab), dtype=tf.int64)
# 触发return_state=True的调用，生成对应签名
_, _ = model(test_input, return_state=True)
# 保存模型
model.save('quoteModel')

方法三：包装加载后的模型适配输入

如果不想修改训练代码，可创建包装类统一处理输入形状：

# 加载原始模型后创建包装类
class WrappedRNNModel(tf.keras.Model):
    def __init__(self, original_model):
        super().__init__()
        self.original_model = original_model
    
    def call(self, inputs, states=None, return_state=False, training=False):
        # 调整输入到(None,100)的形状
        seq_len = tf.shape(inputs)[1]
        if seq_len < 100:
            inputs = tf.pad(inputs, [[0,0], [0, 100 - seq_len]])
        inputs = tf.slice(inputs, [0,0], [-1, 100])
        return self.original_model(inputs, states=states, return_state=return_state, training=training)

# 替换原加载逻辑
new_model = WrappedRNNModel(tf.keras.models.load_model('quoteModel'))

3. 额外检查点 #

• 确认模型call方法的参数顺序为inputs, states=None, return_state=False, training=False，与调用时的关键字参数对应。
• 若训练时计划支持可变长度输入，需将模型输入层的input_shape设置为(None,)而非固定的(100,)。

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