使用 TensorFlow Serving 和 Docker 快速提供 ML 服务
2018 年 11 月 02 日
由 Google Brain 团队的技术项目经理 Gautam Vasudevan 和软件工程师 Abhijit Karmarkar 发布

快速轻松地提供机器学习模型是将实验从实验转移到生产中的关键挑战之一。提供机器学习模型是指将训练好的模型投入使用,使其能够提供预测请求。在生产环境中提供服务时,您需要确保环境可重现、强制隔离并安全。为此,最简单的方法之一是使用 TensorFlow ServingDocker 提供机器学习模型。Docker 是一种将软件打包成名为容器的单元的工具,容器包含运行软件所需的一切。
在 Docker 容器中运行的 TensorFlow Serving
自从 TensorFlow Serving 1.8 发布以来,我们一直在改进对 Docker 支持。我们现在为 CPU 和 GPU 模型提供 Docker 镜像 用于提供服务和开发。为了了解使用 TensorFlow Serving 部署模型的简便程度,让我们尝试将 ResNet 模型投入生产。此模型在 ImageNet 数据集上进行训练,并以 JPEG 图像作为输入,返回图像的分类类别。

我们的示例假设您运行的是 Linux,但它也应该在 macOS 或 Windows 上进行少量修改后即可运行。

使用 TensorFlow Serving 和 Docker 提供 ResNet 服务

第一步是 安装 Docker CE。这将为您提供运行和管理 Docker 容器所需的所有工具。

TensorFlow Serving 使用 SavedModel 格式作为其 ML 模型。SavedModel 是一种与语言无关、可恢复、独立的序列化格式,使高级系统和工具能够生成、使用和转换 TensorFlow 模型。有几种方法可以 导出 SavedModel(包括 从 Keras 导出)。在本练习中,我们将直接下载 预训练的 ResNet SavedModel
$ mkdir /tmp/resnet
$ curl -s https://storage.googleapis.com/download.tensorflow.org/models/official/20181001_resnet/savedmodels/resnet_v2_fp32_savedmodel_NHWC_jpg.tar.gz | tar --strip-components=2 -C /tmp/resnet -xvz
我们现在应该在 /tmp/resnet 中有一个包含我们模型的文件夹。我们可以通过运行以下命令进行验证
$ ls /tmp/resnet
1538687457
现在我们有了模型,使用 Docker 提供服务就像拉取最新的发布的 TensorFlow Serving 服务环境镜像并将其指向模型一样简单
$ docker pull tensorflow/serving
$ docker run -p 8501:8501 --name tfserving_resnet \
--mount type=bind,source=/tmp/resnet,target=/models/resnet \
-e MODEL_NAME=resnet -t tensorflow/serving &
…
… main.cc:327] Running ModelServer at 0.0.0.0:8500…
… main.cc:337] Exporting HTTP/REST API at:localhost:8501 …
分解命令行参数,我们有
  • -p 8501:8501 : 将容器的端口 8501(TF Serving 响应 REST API 请求的位置)发布到主机的端口 8501
  • --name tfserving_resnet : 为我们正在创建的容器命名为“tfserving_resnet”,以便我们以后可以引用它
  • --mount type=bind,source=/tmp/resnet,target=/models/resnet : 将主机的本地目录 (/tmp/resnet) 挂载到容器 (/models/resnet) 上,以便 TF Serving 可以从容器内部读取模型。
  • -e MODEL_NAME=resnet : 告诉 TensorFlow Serving 加载名为“resnet”的模型
  • -t tensorflow/serving : 运行基于服务镜像“tensorflow/serving”的 Docker 容器
接下来,让我们下载 python 客户端脚本,它将发送服务模型图像并获取预测结果。我们还将测量服务器响应时间。
$ curl -o /tmp/resnet/resnet_client.py 
https://raw.githubusercontent.com/tensorflow/serving/master/tensorflow_serving/example/resnet_client.py
此脚本将下载一张猫的图像,并将其重复发送到服务器,同时测量响应时间,如脚本的主循环所示
# The server URL specifies the endpoint of your server running the ResNet
# model with the name "resnet" and using the predict interface.
SERVER_URL = 'http://localhost:8501/v1/models/resnet:predict'

...

# Send few actual requests and time average latency.                                                                                                                                                                   
total_time = 0
num_requests = 10
for _ in xrange(num_requests):
    response = requests.post(SERVER_URL, data=predict_request)
response.raise_for_status()
total_time += response.elapsed.total_seconds()
prediction = response.json()['predictions'][0]

print('Prediction class: {}, avg latency: {} ms'.format(
prediction['classes'], (total_time*1000)/num_requests))
此脚本使用 requests 模块,因此如果您还没有安装它,则需要 安装它。通过运行此脚本,您应该看到类似于以下内容的输出
$ python /tmp/resnet/resnet_client.py
Prediction class: 282, avg latency: 185.644 ms
如您所见,使用 TensorFlow Serving 和 Docker 启动模型非常简单。您甚至可以 创建您自己的自定义 Docker 镜像,其中嵌入您的模型,以便更轻松地部署。

通过构建优化的服务二进制文件来提高性能

现在我们有一个在 Docker 中提供服务的模型,您可能已经注意到 TensorFlow Serving 的日志消息看起来像
Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
发布的 TensorFlow Serving Docker 镜像旨在适用于尽可能多的 CPU 架构,因此为了最大限度地提高兼容性,会省略一些优化。如果您没有看到此消息,则您的二进制文件可能已经针对您的 CPU 进行了优化。

根据您的模型执行的操作,这些优化可能会对您的服务性能产生重大影响。幸运的是,将您自己的优化服务镜像组合在一起非常简单。

首先,我们要构建一个优化的 TensorFlow Serving 版本。最简单的方法是构建官方的 Tensorflow Serving 开发环境 Docker 镜像。它有一个很好的特性,可以自动生成针对构建镜像的系统的 优化 TensorFlow Serving 二进制文件。为了区分我们创建的镜像和官方镜像,我们将为镜像名称添加前缀 $USER/。我们将此构建的开发镜像称为 $USER/tensorflow-serving-devel
$ docker build -t $USER/tensorflow-serving-devel \
-f Dockerfile.devel \ 
https://github.com/tensorflow/serving.git#:tensorflow_serving/tools/docker
构建 TensorFlow Serving 开发镜像可能需要一段时间,具体取决于您的机器速度。完成后,让我们构建一个新的服务镜像,其中包含我们优化的二进制文件,并将其命名为 $USER/tensorflow-serving
$ docker build -t $USER/tensorflow-serving \
--build-arg TF_SERVING_BUILD_IMAGE=$USER/tensorflow-serving-devel \ https://github.com/tensorflow/serving.git#:tensorflow_serving/tools/docker
现在我们有了新的服务镜像,让我们重新启动服务器
$ docker kill tfserving_resnet
$ docker run -p 8501:8501 --name tfserving_resnet \
  --mount type=bind,source=/tmp/resnet,target=/models/resnet \
  -e MODEL_NAME=resnet -t $USER/tensorflow-serving &
最后运行我们的客户端
$ python /tmp/resnet/resnet_client.py
Prediction class: 282, avg latency: 84.8849 ms
在我们的机器上,我们看到每个预测平均加速超过 100 毫秒(119%),使用了我们本地的优化二进制文件。根据您的机器(和模型),您可能会看到不同的结果。

最后,请随时停止 TensorFlow Serving 容器
$ docker kill tfserving_resnet
现在您已经使用 Docker 运行 TensorFlow Serving,您可以轻松地将机器学习模型部署到容器中,同时最大限度地提高部署和性能的简便性。

请阅读我们的 通过 Docker 使用 TensorFlow Serving 文档以获取更多详细信息,并关注我们的 GitHub 项目 以了解最新信息。
下一篇文章
Serving ML Quickly with TensorFlow Serving and Docker

由 Google Brain 团队的技术项目经理 Gautam Vasudevan 和软件工程师 Abhijit Karmarkar 发布

快速轻松地提供机器学习模型是将实验从实验转移到生产中的关键挑战之一。提供机器学习模型是指将训练好的模型投入使用,使其能够提供预测请求。在生产环境中提供服务时,…