TensorBoard What-If Tool对接TensorFlow Serving模型报错及配置求助

问题背景 #

• 数据集：1万行、370列CSV，多数为数值型列+唯一ID列，目标列为整数型
• 已完成：按官方流程构建模型，生成saved_model.pb并通过Docker部署TensorFlow Serving
• 遇到的问题：WIT中触发“Expected one input Tensor”错误，不清楚如何配置输入输出及添加classify/regress/predict签名，需同时支持分类、回归功能

1. 重构模型输入，适配单张量要求 #

WIT期望模型仅接收一个输入张量，所以需要把所有特征列（除ID和目标列）合并成一个二维张量输入：

import pandas as pd
import tensorflow as tf
from tensorflow.keras import layers

# 加载并预处理数据
df = pd.read_csv("your_data.csv")
id_col = df["id_column"]
target_col = df["target_column"]
# 移除ID和目标列，保留所有特征
features = df.drop(["id_column", "target_column"], axis=1)

# 转换为张量格式
X = tf.convert_to_tensor(features.values, dtype=tf.float32)
y = tf.convert_to_tensor(target_col.values, dtype=tf.int32)
num_classes = target_col.nunique()  # 自动获取分类类别数

# 构建多输出模型（同时支持分类、回归）
model = tf.keras.Sequential([
    layers.Dense(64, activation='relu', input_shape=(features.shape[1],)),
    layers.Dense(32, activation='relu'),
    # 分类输出层：多分类用softmax，二分类可换sigmoid
    layers.Dense(num_classes, activation='softmax', name='classification_output'),
    # 回归输出层：直接输出数值
    layers.Dense(1, name='regression_output')
])

# 编译模型，同时指定两个任务的损失函数
model.compile(
    optimizer='adam',
    loss={
        'classification_output': 'sparse_categorical_crossentropy',
        'regression_output': 'mean_squared_error'
    },
    metrics={
        'classification_output': 'accuracy',
        'regression_output': 'mae'
    }
)

# 训练模型
model.fit(X, {'classification_output': y, 'regression_output': y}, epochs=10, batch_size=32)

2. 添加TensorFlow Serving要求的签名 #

必须为classify/regress/predict三个API定义符合规范的签名函数：

# 定义分类签名函数（适配Classify API格式）
@tf.function(input_signature=[tf.TensorSpec(shape=(None, features.shape[1]), dtype=tf.float32, name='inputs')])
def classify_fn(inputs):
    pred_scores = model(inputs)['classification_output']
    # Classify API要求classes为字符串类型，把整数类别转成字符串
    pred_classes = tf.strings.as_string(tf.argmax(pred_scores, axis=1))
    return {'scores': pred_scores, 'classes': pred_classes}

# 定义回归签名函数（适配Regress API格式）
@tf.function(input_signature=[tf.TensorSpec(shape=(None, features.shape[1]), dtype=tf.float32, name='inputs')])
def regress_fn(inputs):
    pred_values = model(inputs)['regression_output']
    return {'predictions': pred_values}

# 定义通用预测签名函数
@tf.function(input_signature=[tf.TensorSpec(shape=(None, features.shape[1]), dtype=tf.float32, name='inputs')])
def predict_fn(inputs):
    return model(inputs)

# 保存模型并绑定签名
save_dir = './saved_model'
tf.saved_model.save(
    model,
    save_dir,
    signatures={
        'serving_default': predict_fn,
        'classify': classify_fn,
        'regress': regress_fn
    }
)

3. 验证签名正确性 #

用saved_model_cli工具检查签名是否符合要求：

saved_model_cli show --dir ./saved_model --all

输出中应能看到classify、regress、serving_default三个签名，每个签名的输入输出形状、类型均需匹配预期。

4. 重新部署TensorFlow Serving #

替换旧模型，用Docker重新启动服务：

docker run -p 8501:8501 -v "$(pwd)/saved_model:/models/your_model" -e MODEL_NAME=your_model tensorflow/serving

5. 在What-If Tool中配置使用 #

1. 打开TensorBoard进入WIT界面，选择“TensorFlow Serving”作为模型来源
2. 输入模型地址：http://localhost:8501/v1/models/your_model
3. 分类任务选“Classify API”，回归任务选“Regress API”
4. 加载CSV时，确保特征列顺序和模型输入的特征顺序完全一致；将ID列设置为WIT的“ID列”用于标识样本

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