什么是TensorFlow serving？

TensorFlow Serving (TFS) 是Google开发的一款高性能、灵活的系统，专门用于将机器学习模型（尤其是TensorFlow模型）投入生产环境。它解决了模型版本管理、高并发请求处理和硬件加速等关键部署问题。

本文将聚焦如何使用Docker部署一个简单的Keras模型，并通过RESTful API和gRPC进行推理调用。

什么是TensorFlow Serving？

TFS作为模型与外部应用之间的桥梁，提供了标准的API接口，允许客户端无需关心底层模型加载、卸载和版本切换的复杂性。其核心优势在于：

1. 模型版本控制： 自动处理新旧模型的切换，实现零停机更新。
2. 高性能： 基于C++编写，支持批量处理和各种硬件加速（如GPU）。
3. 标准化接口： 提供RESTful API（HTTP/JSON）和gRPC API（高性能二进制协议）两种标准接口。

步骤一：准备SavedModel格式的模型

TensorFlow Serving要求模型必须以SavedModel格式保存，并且遵循特定的目录结构：///。

以下是如何保存一个简单的Keras模型的示例：

import tensorflow as tf
import numpy as np
import os

# 定义模型输出路径和版本号
MODEL_DIR = './tf_serving_model'
version = 1
saved_model_path = os.path.join(MODEL_DIR, str(version))

# 1. 创建一个简单的Keras模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(5,)),
    tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy')

# 2. 模拟训练并保存模型
# (通常生产中是加载训练好的模型)
print(f"Saving model to {saved_model_path}")
tf.saved_model.save(model, saved_model_path)

# 验证目录结构
# 此时应存在：./tf_serving_model/1/variables/ 和 ./tf_serving_model/1/assets/

步骤二：使用Docker部署TensorFlow Serving

利用官方Docker镜像可以最快速度启动TFS服务。我们使用卷挂载将本地的模型目录映射到容器内部。

请确保将上一步生成的tf_serving_model目录路径替换为实际的本地绝对路径。

# 假设模型目录在当前工作目录下
MODEL_BASE_PATH=$(pwd)/tf_serving_model
MODEL_NAME="my_simple_classifier"

docker run -d --rm \
    -p 8501:8501 -p 8500:8500 \
    --mount type=bind,source=$MODEL_BASE_PATH,target=/models/$MODEL_NAME \
    -e MODEL_NAME=$MODEL_NAME \
    tensorflow/serving:latest

# 端口说明：
# 8501: gRPC 端口 (高性能)
# 8500: RESTful API 端口 (易于调试)

启动成功后，服务日志会显示模型已成功加载。

步骤三：通过RESTful API进行推理调用

RESTful API通过HTTP POST请求，以JSON格式发送数据。这是最简单、最通用的测试和集成方式。

API Endpoint: http://localhost:8501/v1/models/{model_name}:predict

import requests
import json
import numpy as np

# 模拟输入数据：一个batch，每个样本有5个特征
input_data = np.array([[0.1, 0.2, 0.3, 0.4, 0.5], [1.1, 1.2, 1.3, 1.4, 1.5]])

# 构造JSON请求体
# 注意：键名必须是 'instances' 或者 SavedModel 定义的签名输入名称（默认为 'input_1' 或 'serving_default'）
request_body = {
    "instances": input_data.tolist()
}

# 发送请求
REST_URL = "http://localhost:8501/v1/models/my_simple_classifier:predict"
headers = {"content-type": "application/json"}

response = requests.post(REST_URL, data=json.dumps(request_body), headers=headers)

# 打印结果
print(f"HTTP Status Code: {response.status_code}")
if response.status_code == 200:
    predictions = response.json()['predictions']
    print("Predictions:")
    print(np.array(predictions))
else:
    print(f"Error: {response.text}")

步骤四：通过gRPC进行推理调用（高性能）

对于需要极低延迟和高吞吐量的场景，gRPC是首选。它使用Protocol Buffers进行序列化，效率远高于JSON。

虽然gRPC设置稍微复杂，需要安装grpcio和tensorflow-serving-api，但其性能提升是巨大的。

首先安装必要的库：

pip install grpcio tensorflow-serving-api numpy

Python gRPC客户端示例：

import grpc
import numpy as np
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
from tensorflow.com.google.protobuf import wrappers_pb2
from tensorflow.make_tensor_proto import make_tensor_proto

# 1. 建立通道
channel = grpc.insecure_channel('localhost:8500') # 注意gRPC使用8500端口
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

# 2. 构造请求
request = predict_pb2.PredictRequest()
request.model_spec.name = 'my_simple_classifier'
request.model_spec.signature_name = 'serving_default'

# 3. 准备数据 (与RESTful API使用相同的数据)
input_data = np.array([[0.1, 0.2, 0.3, 0.4, 0.5], [1.1, 1.2, 1.3, 1.4, 1.5]], dtype=np.float32)

# 将numpy array转换为TensorProto格式
request.inputs['dense_input'].CopyFrom(
    make_tensor_proto(input_data, shape=input_data.shape)
)

# 4. 发送请求
try:
    result = stub.Predict(request, timeout=5.0) # 5秒超时
    # 解析结果
    output_tensor = result.outputs['dense_1'] # 模型的输出层名称
    predictions = np.array(output_tensor.float_val).reshape(input_data.shape[0], -1)
    print("gRPC Predictions:")
    print(predictions)
except grpc.RpcError as e:
    print(f"gRPC Error: {e.details()}")

总结

TensorFlow Serving是部署TensorFlow模型不可或缺的工具。通过使用标准化的SavedModel格式和Docker容器化部署，我们可以快速构建一个高性能、可扩展的推理服务。对于生产环境，建议优先使用gRPC接口以获得最佳的性能体验。