TensorFlow v2.10修改签名函数输出为标签-结果映射适配TF Serving

核心问题分析 #

1. 语法错误：原代码中return语句后直接换行，导致实际返回None，触发后续签名验证错误。
2. 输出结构不符合要求：TF SavedModel签名要求输出必须是Tensor、Tensor序列或字符串到Tensor的字典，原代码返回的嵌套Python字典包含非Tensor值（格式化字符串），不符合规范。
3. 模型保存对象错误：保存时传入的是tf_model_wrapper.model（原始Keras模型），而非TFModel实例，导致无法扫描到自定义的prediction签名，自动签名机制失效。

修改后的完整代码 #

import tensorflow as tf

class TFModel(tf.Module):
    def __init__(self, model: tf.keras.Model) -> None:
        self.labels = ['label1', 'label2']
        self.model = model
            
    @tf.function(input_signature=[tf.TensorSpec(shape=(None, ), dtype=tf.string)])
    def prediction(self, pagetext: tf.Tensor):
        # --- 补充你的文本预处理逻辑 ---
        # 示例：假设需要将字符串转为模型可接受的输入张量
        # processed_input = 你的预处理代码（如分词、转embedding等）
        # --- 执行模型推理 ---
        results = self.model(processed_input)  # 输出shape: (batch_size, 2)
        
        # 保留5位小数，避免使用Python字符串格式化（保持Tensor类型）
        rounded_probs = tf.round(results * 100000) / 100000
        
        # 构建符合要求的输出结构
        batch_output = []
        for idx in range(tf.shape(rounded_probs)[0]):
            sample_probs = rounded_probs[idx]
            batch_output.append({
                self.labels[0]: sample_probs[0],
                self.labels[1]: sample_probs[1]
            })
        
        return {'results': batch_output}

# 保存模型（注意保存TFModel实例，而非内部的Keras模型）
tf_model_wrapper = TFModel(classifier_model)
tf.saved_model.save(tf_model_wrapper, saved_model_path)

关键修改说明 #

• 修复return语法：确保返回值与return语句在逻辑上同属一行（或用括号包裹多行结构），避免返回空值。
• 保持Tensor类型输出：用tf.round处理概率值的精度，避免将Tensor转为Python字符串，确保输出符合签名要求。
• 正确设置输出结构：返回的results是一个包含字典的列表，每个字典的键为自定义标签，值为对应概率的Tensor，完全符合TF Serving的输出规范。
• 保存正确对象：保存TFModel实例而非原始Keras模型，让TF自动扫描并设置自定义的prediction方法为serving_default签名。

验证输出 #

通过TF Serving调用时，会自动将Tensor序列化为你预期的JSON格式：

{
   "results": [
      { "label1": 0.12345, "label2": 0.87655 },
      { "label1": 0.98765, "label2": 0.01235 }
   ]
}

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