TensorFlow Serving是Google开发的一款高效云端TensorFlow模型托管框架。使用TensorFlow训练完成的模型可以直接导出为 saved_model 格式,然后直接用在TensorFlow Serving中,并自动实现batching。TensorFlow对外提供gRPC与REST两种调用调用接口以方便调用,对于资源有限的小型公司,是一个不错的选择。

封面图像:Bosco Shots on Unsplash

背景

在TensorFlow官网可以找到关于TensorFlow Serving(以下简称Serving)的详细介绍。官方提供了APT与Docker两种安装方式。我在实际测试时发现这两种方案在高并发条件下会出现奔溃现象。为了排除问题,有可能需要从源码编译。

编译对象

Serving面向的是最终的应用部署,该环境下使用GPU以加速运算是大概率事件。因此编译对象为TensorFlow Serving GPU版本。

编译环境

Ubuntu 16.04

编译过程

整个过程大致分为:

  1. 安装依赖项。
  2. 设定编译条件。
  3. 完成编译过程。

让我们开始吧!

安装依赖项

Serving的主要依赖项为Nvidia CUDA、CUDNN与Bazel。

CUDA

首先在官方网站下载CUDA安装包:

https://developer.nvidia.com/cuda-toolkit

推荐使用CUDA 10。注意下载适用Ubuntu 16.04的deb包,如果要离线安装的话选择 local 选项。

下载完成后,按照官方安装说明完成CUDA安装:

https://developer.download.nvidia.com/compute/cuda/10.0/Prod/docs/sidebar/CUDA_Installation_Guide_Linux.pdf

cuDNN

同样在官方下载cuDNN安装包:

https://developer.nvidia.com/cudnn

下载时要求登录,注册后即可。

下载时留意下版本号,要与CUDA版本对应起来。对于编译来说需要安装runtime与developer版。

下载完成后,参照官方说明完成cuDNN安装:

https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html

Bazel

Bazel是Google开发的构建系统。它的安装较为简单,留意一下TensorFlow对Bazel的版本要求即可。

首先从官方页面下载release版本:

https://github.com/bazelbuild/bazel/releases

官方提供了deb包,所以 dpkg -i 应该就能搞定。

设定编译条件

Serving的编译完全可以参照官方的Docker file来实现。经过测试,编译GPU版本的Serving需要设定系统环境变量:

export TF_NEED_CUDA=1
export TF_CUDA_VERSION=10.0
export TF_CUDNN_VERSION=7

为了安全起见,你也可以参照Docker file为对应的so库文件设定软链接。

完成编译过程

以上步骤完成后,开始编译:

bazel build --config=cuda --config=nativeopt --copt="-fPIC" tensorflow_serving/model_servers:tensorflow_model_server

编译过程中Bazel会下载包括TensorFlow在内一众依赖项源码。因此具体编译时长与网络与运算性能有关。

编译完成后,查看版本:

bazel-bin/tensorflow_serving/model_servers/tensorflow_model_server --version
TensorFlow ModelServer: 0.0.0+dev.sha.66e40b0
TensorFlow Library: 1.12.0

OK,编译完成。