5.3 TensorFlow Serving 模型部署

5.3 TensorFlow Serving 模型部署

TensorFlow Serving 模型部署详解

1. TensorFlow Serving 概述

TensorFlow Serving 的核心思想是将模型视为服务，通过 gRPC 或 RESTful API 提供预测功能。其主要优点包括：

• 易于部署： 简化了模型部署流程，无需复杂的配置。

• 版本控制： 支持多个模型版本同时存在，方便回滚和 A/B 测试。

• 高性能： 采用高性能的 gRPC 协议，提供低延迟的预测服务。

• 动态更新： 可以在不中断服务的情况下更新模型。

2. 模型准备

在部署模型之前，需要将其保存为 TensorFlow Serving 可以识别的格式。通常，我们使用 tf.saved_model.save 函数将模型保存为 SavedModel 格式。

import tensorflow as tf
# 创建一个简单的模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dense(10)
])
# 定义优化器和损失函数
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
# 创建一些虚拟数据
import numpy as np
data = np.random.rand(100, 784).astype(np.float32)
labels = np.random.randint(0, 10, 100)
# 训练模型
model.fit(data, labels, epochs=5)
# 保存模型
version = 1  # 模型版本号
export_path = f'./model/mnist/{version}'  # 模型保存路径
tf.saved_model.save(model, export_path)
​

代码解释：

1. 创建模型： 创建了一个简单的 MNIST 分类模型。

2. 编译模型： 定义了优化器、损失函数和评估指标。

3. 训练模型： 使用虚拟数据训练模型。

4. 保存模型： 使用 tf.saved_model.save 函数将模型保存到指定路径。export_path 包含模型的基本路径和版本号。

SavedModel 目录结构：

SavedModel 目录结构如下：

mnist/
  1/
    assets/
    variables/
      variables.data-00000-of-00001
      variables.index
    saved_model.pb
​

• assets/: 存放模型需要的额外资源文件。

• variables/: 存放模型的权重和偏置等变量。

• saved_model.pb: 存放模型的图结构和其他元数据。

3. 安装 TensorFlow Serving

可以使用 Docker 或直接安装 TensorFlow Serving。这里以 Docker 为例：

docker pull tensorflow/serving
​

4. 启动 TensorFlow Serving

使用以下命令启动 TensorFlow Serving：

docker run -t --rm -p 8501:8501 \
    -v "$(pwd)/model:/models" \
    -e MODEL_NAME=mnist \
    tensorflow/serving
​

命令解释：

• -t --rm: 在终端中运行容器，并在容器停止时自动删除。

• -p 8501:8501: 将主机的 8501 端口映射到容器的 8501 端口（gRPC 默认端口）。

• -v "$(pwd)/model:/models": 将主机上的模型目录挂载到容器的 /models 目录。

• -e MODEL_NAME=mnist: 设置模型名称为 mnist。

流程图：

5. 发送预测请求

可以使用 gRPC 或 REST API 向 TensorFlow Serving 发送预测请求。

5.1 使用 gRPC 发送请求

首先，需要安装 grpcio 和 tensorflow-serving-api：

pip install grpcio tensorflow-serving-api
​

然后，编写客户端代码：

import grpc
import tensorflow as tf
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
# 连接到 TensorFlow Serving
channel = grpc.insecure_channel('localhost:8501')
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
# 创建预测请求
request = predict_pb2.PredictRequest()
request.model_spec.name = 'mnist'
request.model_spec.signature_name = 'serving_default'
# 构造输入数据
data = np.random.rand(1, 784).astype(np.float32)
request.inputs['dense_input'].CopyFrom(
    tf.compat.v1.make_tensor_proto(data)
)
# 发送请求并获取响应
result = stub.Predict(request, timeout=10.0)
# 处理响应
print(result)
​

代码解释：

1. 连接到 TensorFlow Serving： 使用 grpc.insecure_channel 创建到 TensorFlow Serving 的连接。

2. 创建预测请求： 创建 PredictRequest 对象，并设置模型名称和签名名称。

3. 构造输入数据： 将输入数据转换为 tf.TensorProto 格式，并将其添加到请求中。

4. 发送请求并获取响应： 使用 stub.Predict 函数发送请求，并获取响应。

5. 处理响应： 从响应中提取预测结果。

5.2 使用 REST API 发送请求

import requests
import json
# 构造请求数据
data = json.dumps({
    "instances": [np.random.rand(784).tolist()]
})
# 发送 POST 请求
headers = {"content-type": "application/json"}
json_response = requests.post('http://localhost:8501/v1/models/mnist:predict', data=data, headers=headers)
# 处理响应
predictions = json.loads(json_response.text)['predictions']
print(predictions)
​

代码解释：

1. 构造请求数据： 将输入数据转换为 JSON 格式。

2. 发送 POST 请求： 使用 requests.post 函数发送 POST 请求到 TensorFlow Serving 的 REST API 端点。

3. 处理响应： 从响应中提取预测结果。

6. 模型版本管理

TensorFlow Serving 支持多个模型版本同时存在。可以将不同版本的模型保存在不同的子目录中，TensorFlow Serving 会自动加载最新版本。

例如，可以创建以下目录结构：

mnist/
  1/
    ...
  2/
    ...
​

TensorFlow Serving 会自动加载版本号最高的模型（在本例中为版本 2）。

7. 模型更新

可以通过以下步骤更新模型：

1. 保存新版本的模型： 将新版本的模型保存到新的子目录中。

2. TensorFlow Serving 自动加载新模型： TensorFlow Serving 会定期检查模型目录，并自动加载新版本的模型。

可以使用 --model_config_file 参数指定模型配置，从而更精细地控制模型的加载和卸载。

8. 高级配置

TensorFlow Serving 提供了许多高级配置选项，例如：

• 模型配置： 可以使用模型配置文件来指定模型的加载策略、资源限制等。

• 签名定义： 可以定义不同的签名来支持不同的输入输出格式。

• 批处理： 可以启用批处理来提高吞吐量。

9. 总结

本文详细介绍了 TensorFlow Serving 的模型部署过程，包括模型准备、TensorFlow Serving 安装启动、发送预测请求、模型版本管理和模型更新。通过本文的学习，您可以快速上手 TensorFlow Serving，并将其应用到实际的生产环境中。