如何在tf serving中集成预处理并获取分类类别与预测概率

问题背景 #

开展文本分类任务时，需对输入自由文本（例如"Today the weather is good"）预测所属类别，计划使用TF Serving部署推理服务。当前本地推理流程不会直接将原始自由文本输入模型，需先执行一系列预处理操作：

• 文本转小写
• 调用remove_URL移除链接
• 调用remove_punct移除标点
• 调用remove_stopwords移除停用词
• 通过tokenizer.texts_to_sequences转换为序列
• 调用pad_sequences做等长padding，最终得到形如以下结构的x_ready特征矩阵输入模型：

[[
  0 0 0 0 
  0 0 0 0 
  1 5 987 54
]]

本地单样本预测代码如下，推理后会输出预测概率与对应类别：

def predict_single(x):    
    # 文本转小写
    x = x.lower()
    x = remove_URL(x)
    x = remove_punct(x)
    x = remove_stopwords(x)
    
    x = tokenizer.texts_to_sequences([x])
    # 等长填充
    x_ready = pad_sequences(x, maxlen=maxlen)
    print(x)
    pred = model.predict(x_ready)
    print("Probability", pred.max(axis=1))
    pred = pred.argmax(axis=1)
    pred = le.classes_[pred]
    
    return pred[0]
    
    # 清洗后文本
    return x

本地运行输出示例：

[OUT]
Probability [0.7528353]
Class 3

目前已通过Docker拉取tensorflow/serving镜像启动容器，执行以下命令启动TF Serving服务：

docker pull tensorflow/serving

docker run -it -v C:\Code\tensorflow\demo:/demo -p 8601:8601 --entrypoint /bin/bash tensorflow/serving

tensorflow_model_server --rest_api_port=8601 --model_name=model --model_base_path=/demo/models/

模型训练完成后通过以下方式保存：

model.fit(...)
model.save("models/1")
model.save("models/2")

需解决两个核心问题：

1. 如何将上述文本预处理逻辑集成到TF Serving的推理流程中，实现传入原始文本即可完成推理，无需在服务外额外做预处理
2. 调用TF Serving接口时，如何同时获取返回的预测分类类别与对应概率结果

解决方案 #

预处理逻辑集成到TF Serving #

TF Serving仅支持加载TensorFlow计算图内的逻辑，需把所有预处理逻辑封装为模型的一部分，导出为标准SavedModel格式，不要单独导出纯分类模型，具体操作步骤：

• 将自定义文本预处理逻辑转换为TF原生算子
  本地使用的Python原生字符串处理、第三方库实现的停用词移除逻辑无法直接序列化进TF计算图，需替换为tf.strings系列API实现：
  • 转小写：直接调用tf.strings.lower
  • 移除URL：用tf.strings.regex_replace匹配URL正则做替换
  • 移除标点：同样用tf.strings.regex_replace匹配标点符号正则做替换
  • 移除停用词：先把停用词表转为TF静态哈希表，用tf.strings.split分词后过滤命中停用词的token再拼接
• 将Tokenizer和Padding逻辑封装进模型
  训练完成后不要直接保存原分类模型，先提取Tokenizer的词表构建tf.keras.layers.TextVectorization层，把训练用的词表、固定序列长度maxlen、padding规则配置到该层，这层可原生完成分词、转序列、等长padding全流程，输出和本地tokenizer.texts_to_sequences+pad_sequences的结果完全对齐。
• 拼接端到端模型后导出
  新建一个以原始字符串为输入的新模型，结构为：输入层（string类型） -> 预处理子模块 -> 原训练好的分类模型，调用model.save()导出整个端到端模型到models目录，TF Serving加载该模型后，直接接收原始文本字符串输入即可完成全流程推理，不需要外部做任何预处理。

注意：不要在TF Serving容器里额外跑Python预处理脚本做前置转发，这种方案性能差、链路长，后续版本迭代很容易出现预处理逻辑和模型版本不匹配的问题，把预处理合入计算图是官方推荐的生产方案。

接口同时返回分类类别和概率 #

默认TF Serving返回模型最后一层的原始输出，也就是分类概率向量，需修改导出模型的输出签名，同时输出两个结果：

• 第一个输出是概率分布向量，用来提取最大概率值
• 第二个输出是经argmax计算得到的类别索引，再把标签编码器le.classes_转为TF静态查找表，把类别索引映射为真实类别名称，作为模型的第二个输出
  导出模型时显式指定输出签名，把pred_prob（概率向量）和pred_label（真实类别名）都注册为模型输出，调用TF Serving接口时就会同时返回这两个字段，直接从返回结果里取最大概率值和对应类别即可，不需要客户端再做索引到类别的映射。

接口调用返回示例结构如下：

{
    "predictions": [
        {
            "pred_prob": [0.12, 0.08, 0.05, 0.75],
            "pred_label": "Class 3"
        }
    ]
}

从pred_prob里取最大值就是对应预测置信度，pred_label就是最终分类结果。

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