如何使用TensorFlow Serving调用模型指定的自定义推理方法？

实现方案很简单，不需要单独创建新的servable，通过给SavedModel注册自定义签名即可实现直接调用query_with_exclusions方法，具体操作如下： #

• 第一步：导出模型时注册自定义方法签名
  在保存SCaNN模型训练完成后、导出SavedModel阶段，通过tf.saved_model.save的signatures参数，将query_with_exclusions方法注册为独立的推理签名，参考代码如下：

import tensorflow as tf
import tensorflow_recommenders as tfrs

# 加载训练好的SCaNN模型实例
scann_model = tfrs.layers.factorized_top_k.ScaNN(...)
# 此处省略模型加载、参数配置逻辑

# 定义方法签名注册
signatures = {
    # 保留默认的call方法签名，不影响原有逻辑使用
    "serving_default": scann_model.call.get_concrete_function(
        tf.TensorSpec([None, 128], tf.float32, name="query_embedding")
    ),
    # 新增query_with_exclusions方法的签名
    "query_with_exclusions": scann_model.query_with_exclusions.get_concrete_function(
        # 此处TensorSpec配置要和方法实际接收的参数维度、数据类型完全匹配
        tf.TensorSpec([None, 128], tf.float32, name="query_embedding"),
        tf.TensorSpec([None, None], tf.int64, name="exclusion_item_ids")
    )
}

# 导出带自定义签名的模型
tf.saved_model.save(scann_model, "./exported_scann_model", signatures=signatures)

• 第二步：调用指定签名的推理接口
  导出的模型可以直接用你当前使用的官方SCaNN镜像部署，不需要修改任何服务端配置。调用REST接口时，在请求体中增加signature_name字段指定为你注册的query_with_exclusions即可，请求示例如下：

{
  "signature_name": "query_with_exclusions",
  "inputs": {
    "query_embedding": [[0.1, 0.3, ..., 0.2]],
    "exclusion_item_ids": [[1001, 1005, 1200]]
  }
}

请求路径保持原有/v1/models/[你的模型名称]:predict不变即可。

这个方案是TensorFlow Serving官方支持的标准多推理方法暴露方案，相比重新封装call方法生成新servable的方式开发成本低很多，也不会影响原有call方法的正常调用。

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