如何让TensorFlow Serving保留单输出的自定义命名返回对象格式

遇到单输出场景下TensorFlow Serving自动剥离输出名称、返回数组而非命名对象的问题，核心是要让模型的输出结构保持字典形式，而非单个张量，具体可通过以下两种方式实现：

1. 直接在Keras模型中返回字典输出 #

修改模型的输出逻辑，把单个输出张量包装成键值对字典。比如原本模型最后返回单个张量：

def call(self, inputs):
    # ... 模型计算逻辑
    return output_tensor

改成返回字典：

def call(self, inputs):
    # ... 模型计算逻辑
    return {"自定义输出名": output_tensor}

这样导出的模型会默认把字典的键作为输出名称，TensorFlow Serving识别后会返回带命名的对象结构。

2. 自定义签名函数导出SavedModel #

如果不想修改模型结构，可以在导出模型时通过自定义签名函数，强制把输出包装成字典：

import tensorflow as tf
from tensorflow.keras import models

# 假设你已经有训练好的单输出模型
model = models.load_model('./your_trained_model')

# 定义服务签名函数，指定输入输出结构
@tf.function(input_signature=[tf.TensorSpec(shape=(None, 输入维度), dtype=tf.float32, name='input')])
def serving_signature(inputs):
    outputs = model(inputs)
    # 包装成字典，键就是你想要的输出名称
    return {"自定义输出名": outputs}

# 保存模型时指定自定义签名
tf.saved_model.save(model, './saved_for_serving', signatures={'serving_default': serving_signature})

验证方法 #

导出模型后，用saved_model_cli查看签名：

saved_model_cli show --dir ./saved_for_serving --all

确认输出部分显示为字典结构（比如Outputs: "自定义输出名"），部署后请求TensorFlow Serving，返回的JSON就会是带命名的对象列表，格式类似：

{"predictions": [{"自定义输出名": [1.23, 4.56, ...]}]}

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