如何通过JSON向tf.Estimator部署的TensorFlow模型发起REST调用

我已构建并训练了一个TensorFlow模型，采用tf.Estimator范式进行部署。我编写了如下服务函数：

def serving_input_fn(params):
    feature_placeholders = { 'inputs' : tf.placeholder(tf.int64, [None], name='inputs') }
    features = { key: tensor for key, tensor in feature_placeholders.items() }
    return tf.estimator.export.ServingInputReceiver(features, feature_placeholders)

现在我希望能够以application/json为内容类型调用该模型，请问该如何实现？

要让你的tf.Estimator模型支持application/json格式的请求，核心是调整serving_input_fn，让它能正确解析JSON格式的输入数据。目前你的函数依赖原始张量占位符，默认只支持TensorFlow的Example格式或纯张量请求，得针对性修改逻辑：

1. 重构Serving Input函数 #

这里提供两种实用的修改方式，适配不同的JSON输入场景：

方式一：用TensorFlow原生工具解析JSON #

如果你的请求是标准的JSON结构（比如{"inputs": [1,2,3,4]}），可以直接用tf.io.parse_json_example来解析，这种方式完全在TensorFlow图内运行，效率更高：

def serving_input_fn(params):
    # 定义接收JSON字符串的占位符
    json_input = tf.placeholder(dtype=tf.string, shape=[None], name='json_input')
    
    # 定义JSON解析规则，匹配你的输入结构
    feature_spec = {
        'inputs': tf.io.FixedLenFeature([], tf.int64, default_value=0)
    }
    # 解析JSON字符串为模型可用的特征张量
    parsed_features = tf.io.parse_json_example(json_input, feature_spec)
    
    # 调整张量形状为模型预期的一维格式
    inputs_tensor = tf.reshape(parsed_features['inputs'], shape=[-1])
    
    # 返回ServingInputReceiver：features是模型接收的张量，receiver_tensors是外部传入的占位符
    return tf.estimator.export.ServingInputReceiver(
        features={'inputs': inputs_tensor},
        receiver_tensors={'json_input': json_input}
    )

方式二：自定义Python逻辑解析复杂JSON #

如果你的JSON结构比较灵活（比如嵌套数组、多字段混合），可以用tf.py_function嵌入Python的JSON解析逻辑，兼容性更强：

import json

def parse_json_to_tensor(json_strings):
    # 用Python解析每个JSON字符串，整理成输入张量
    input_list = []
    for json_str in json_strings.numpy():
        data = json.loads(json_str.decode('utf-8'))
        input_list.extend(data['inputs'])
    return tf.convert_to_tensor(input_list, dtype=tf.int64)

def serving_input_fn(params):
    json_input = tf.placeholder(dtype=tf.string, shape=[None], name='json_input')
    
    # 在TensorFlow图中调用自定义解析函数
    inputs_tensor = tf.py_function(
        func=parse_json_to_tensor,
        inp=[json_input],
        Tout=tf.int64
    )
    # 手动设置张量形状，让模型能识别输入维度
    inputs_tensor.set_shape([None])
    
    return tf.estimator.export.ServingInputReceiver(
        features={'inputs': inputs_tensor},
        receiver_tensors={'json_input': json_input}
    )

2. 导出支持JSON的模型 #

修改完serving_input_fn后，用estimator.export_saved_model导出模型时指定这个函数即可：

estimator.export_saved_model(
    export_dir_base='./saved_json_model',
    serving_input_receiver_fn=serving_input_fn
)

3. 测试JSON请求 #

模型导出后，用TensorFlow Serving部署，就可以发送application/json请求测试了：

用curl测试： #

curl -X POST http://localhost:8501/v1/models/your_model:predict \
  -H "Content-Type: application/json" \
  -d '{"instances": [{"json_input": "{\"inputs\": [10, 20, 30]}"}]}'

用Python requests测试： #

import requests

url = 'http://localhost:8501/v1/models/your_model:predict'
headers = {'Content-Type': 'application/json'}
request_data = {
    "instances": [
        {"json_input": '{"inputs": [10, 20, 30]}'}
    ]
}
response = requests.post(url, headers=headers, json=request_data)
print(response.json())

关键注意点 #

• 如果你的模型支持批量输入，可以在JSON的inputs里传入多组数据，或者调整解析逻辑适配批量维度。
• 使用tf.io.parse_json_example时，要注意它的输入是每个元素为单个JSON字符串的张量，所以请求里的instances要对应每个独立的JSON实例。
• 部署时确保TensorFlow Serving指向了正确的导出模型目录。

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