Tensorflow Serving 模型部署指南

文章目录

由于python的灵活性和完备的生态库，使得其成为实现、验证ML算法的不二之选。但是工业界要将模型部署到生产环境上，需要考略性能问题，就不建议再使用python端的服务。这个从训练到部署的整个流程如下图所示：

基本可以把工作分为三块：

Saver端：模型的离线训练与导出
Serving端：模型加载与在线预测
Client端：构建请求

本文采用 Saver (python) + Serving (tensorflow serving) + Client (Java) 作为解决方案，记录线上模型部署流程。

1. Saver端：模型的离线训练与导出

部署模型第一步是将训练好的整个模型导出为一系列标准格式的文件，然后即可在不同的平台上部署模型文件。TensorFlow 使用 SavedModel（pb文件）这一格式用于模型部署。与Checkpoint 不同，SavedModel 包含了一个 TensorFlow 程序的完整信息：不仅包含参数的权值，还包含计算图。

SavedModel最终保存结果包含两部分saved_model.pb和variables文件夹。

1.1 saved_model 模型保存与载入

saved_model模块主要用于TensorFlow Serving，模型的保存主要基于
tf.saved_model.builder.SavedModelBuilder 方法。

参考博客：TensorFlow saved_model 模块

1.1.1 简单场景：模型保存

最简单的场景，只是保存/载入模型。核心实现代码如下：

builder = tf.saved_model.builder.SavedModelBuilder(saved_model_path)
builder.add_meta_graph_and_variables(
    sess=sess,
    tags=['lstm_saved_model']
)
builder.save()

详细解读一下上述几行核心代码：

首先构造SavedModelBuilder对象，初始化方法只需要传入用于保存模型的目录名，目录不用预先创建。
add_meta_graph_and_variables方法导入graph的信息以及变量，这个方法假设变量都已经初始化好了，对于每个SavedModelBuilder这个方法一定要执行一次用于导入第一个meta graph。
sess参数：传入当前的session，包含了graph的结构与所有变量。
tags参数：是给当前需要保存的meta graph一个标签，标签名可以自定义。在之后载入模型的时候，需要根据这个标签名去查找对应的MetaGraphDef。找不到就会报如RuntimeError: MetaGraphDef associated with tags ‘foo’ could not be found in SavedModel这样的错。标签也可以选用系统定义好的参数，如 tf.saved_model.tag_constants.SERVING 与tf.saved_model.tag_constants.TRAINING。
save方法就是将模型序列化到指定目录底下。

保存好以后到saved_model_dir目录下，会有一个saved_model.pb文件以及variables文件夹。顾名思义，variables保存所有变量，saved_model.pb用于保存模型结构等信息。

1.1.2 简单场景：模型载入

保存完后，可以在本地载入pb模型，简单测试一下模型的是否正确。

meta_graph_def = tf.saved_model.loader.load(
    sess, 
    ['lstm_saved_model'], 
    saved_model_path
)

第一个参数就是当前的session，第二个参数是在保存的时候定义的meta graph的标签，标签一致才能找到对应的meta graph。第三个参数就是模型保存的目录。load完以后，也是从sess对应的graph中获取需要的tensor来inference。

def test_saved_model_pb(saved_model_path):
    
    test_graph = tf.Graph()
    
    with tf.Session(graph=test_graph) as test_sess:
        
        # 调用 tf.saved_model.loader.load 恢复出meta_graph
        meta_graph_def = tf.saved_model.loader.load(
            sess=test_sess, 
            tags=['lstm_saved_model'],  
            export_dir=saved_model_path
        )
        
        # 从sess对应的graph中获取需要的tensor
        inputXX = test_sess.graph.get_tensor_by_name('inputXX:0')       
        y_pred = test_sess.graph.get_tensor_by_name('ypred:0')
        
        # 加载测试数据
        test_X = ...
        
        # 得到测试结果
        test_output = test_sess.run([y_pred], feed_dict = {inputXX: test_X})

1.1.3 使用SignatureDef：模型保存

传统的tensor导入需要用get_tensor_by_name ，这样就需要一一记住构建阶段tensor的name，很麻烦。

在下面的代码中，我们使用了SignatureDef，将输入输出tensor的信息都进行了封装，并且给他们一个自定义的别名。所以在构建模型的阶段，可以随便给tensor命名，只要在保存训练好的模型的时候，在SignatureDef中给出统一的别名即可。这样，输入输出tensor的具体名称已经完全隐藏，从而实现了创建模型与使用模型的解耦。

第一步：构建 Signature，下面给出构建的两种方式：

### 方法一
# 构建两个字典，inputs 和 outputs，把要存入的变量放入其中
inputs = {'input_XX': tf.saved_model.utils.build_tensor_info(inputX)}
outputs = {'pred_YY': tf.saved_model.utils.build_tensor_info(y_pred)}

# 构建 Signature
signature = tf.saved_model.signature_def_utils.build_signature_def(
    inputs=inputs,
    outputs=outputs,
    method_name=signature_key
)

### 方法二
# 构建需要在新会话中恢复的变量的 TensorInfo
x_tensor_info = tf.saved_model.utils.build_tensor_info(inputX)
y_tensor_info = tf.saved_model.utils.build_tensor_info(y_pred)

# 构建 Signature
signature = tf.saved_model.signature_def_utils.build_signature_def(
    inputs = {'input_XX': x_tensor_info},
    outputs = {'pred_YY': y_tensor_info},
    method_name = signature_key
)

第二步：构造SavedModelBuilder对象，导入graph的信息以及变量，并将模型序列化到指定目录底下。核心代码：

signature_key = 'lstm_signature'

# 构造SavedModelBuilder对象
builder = tf.saved_model.builder.SavedModelBuilder(saved_model_path)
# 导入graph的信息以及变量，并以signature的形式添加要存储的变量
builder.add_meta_graph_and_variables(
    sess=sess,
    tags=['lstm_saved_model'], # tags可以自定义，也可以使用预定义值： tags=[tf.saved_model.tag_constants.TRAINING]
    signature_def_map={signature_key: signature} ,
    clear_devices=True
)
# 将模型序列化到指定目录底下
builder.save()

还是考虑基于LSTM的预测模型，完整版代码如下：

def model_train_saved_model_pb(train_X, train_Y, FLAGS, saved_model_path):
    """
    训练模型；基于saved_model方法将模型保存为pb文件形式
    param:
        train_X: 训练数据X，shape=[train_len, seq_len, num_nodes]
        train_Y: 训练数据Y，shape=[train_len, pre_len, num_nodes]
        FLAGS: 命令行参数
        saved_model_path: 基于saved_model方法所pb文件的路径；可以是一个不存在的目录
    """
    
    # parameters
    seq_len = FLAGS.seq_len
    pre_len = FLAGS.pre_len
    batch_size = FLAGS.batch_size
    lr = FLAGS.learning_rate
    training_epoch = FLAGS.training_epoch
    num_units = FLAGS.num_units
    
    # placeholders
    inputX = tf.placeholder(tf.float32, shape=[None, seq_len, num_nodes], name='inputXX') 
    labelY = tf.placeholder(tf.float32, shape=[None, pre_len, num_nodes], name='labelYY') 
    
    weight_out = {'out': tf.Variable(tf.random_normal([num_units, pre_len], mean=1.0), name='weight_out')}
    bias_out  = {'out': tf.Variable(tf.random_normal([pre_len]), name='bias_out')}
    
    # traffic prediction model
    y_pred, _, _ = lstm_prediction(inputX, weight_out, bias_out) #  y_pred.name = 'ypred'

    
    # model loss
    lambda_loss = 0.0015
    label = tf.reshape(labelY, [-1, num_nodes])
    loss = tf.reduce_mean(tf.nn.l2_loss(y_pred-label), name='loss')  
    
    # error
    error = tf.sqrt(tf.reduce_mean(tf.square(y_pred-label)), name='error') 
    
    # optimizer
    optimizer = tf.train.AdamOptimizer(lr).minimize(loss) # ADAM

    # saver 
    saver = tf.train.Saver(tf.global_variables())  

    # start session 
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        
        batch_loss, batch_rmse = [], []
        time_start = time.time()
        
        for epoch in range(training_epoch):   
            for m in range(total_batch):
                mini_batch = train_X[m*batch_size : (m+1)*batch_size]
                mini_label = train_Y[m*batch_size : (m+1)*batch_size]
                
                _, loss1, rmse1, train_output = sess.run([optimizer, loss, error, y_pred], 
                                                         feed_dict={inputX: mini_batch, labelY: mini_label})    
                batch_loss.append(loss1)
                batch_rmse.append(rmse1) 
            print('Iter:{}'.format(epoch),  
                  'train_loss:{:.4}'.format(batch_loss[-1]),
                  'train_rmse:{:.4}'.format(batch_rmse[-1])) 
   
        # saved_model 开始，定义 signature_key
        signature_key = 'lstm_signature'
        
        
#         # 构建需要在新会话中恢复的变量的 TensorInfo
#         x_tensor_info = tf.saved_model.utils.build_tensor_info(inputX)
#         y_tensor_info = tf.saved_model.utils.build_tensor_info(y_pred)
        
#         # 构建 Signature
#         signature = tf.saved_model.signature_def_utils.build_signature_def(
#             inputs = {'input_XX': x_tensor_info},
#             outputs = {'pred_YY': y_tensor_info},
#             method_name = signature_key
#         )        
        
        # 构建两个字典，inputs 和 outputs，把要存入的变量放入其中
        inputs = {'input_XX': tf.saved_model.utils.build_tensor_info(inputX)} # 输入数据inputX
        outputs = {'pred_YY': tf.saved_model.utils.build_tensor_info(y_pred)} # 预测结果y_pred
        
        # 构建 Signature
        signature = tf.saved_model.signature_def_utils.build_signature_def(
            inputs=inputs,
            outputs=outputs,
            method_name=signature_key
        )
        
        # 建立SavedModelBuilder存储模型，并以signature的形式添加要存储的变量
        builder = tf.saved_model.builder.SavedModelBuilder(saved_model_path)
        builder.add_meta_graph_and_variables(
            sess=sess,
            tags=['lstm_saved_model'], 
            # tags=[tf.saved_model.tag_constants.TRAINING],
            # tags=[tf.saved_model.tag_constants.SERVING],
            signature_def_map={signature_key: signature} ,
            clear_devices=True
        )
        
        # 将 MetaGraphDef 写入磁盘
        builder.save()
        
        time_end = time.time()
        print('Training time:{:.4}'.format(time_end-time_start),'s')

[补充说明]：

1. tags的选择

#  自定义tags
tags=['lstm_saved_model']

# 预定义值TRAINING
tags=[tf.saved_model.tag_constants.TRAINING]

# 预定义值SERVING,若后续放在TFServing上,则必须选择SERVING
tags=[tf.saved_model.tag_constants.SERVING]

注:
(1) tags只可以选一个值, 如果同时写上 TRAINING 和 SERVING 则会报错.
(2) 固化后的模型后续如果部署在TFServing上, 则此处tags必须选择SERVING, 不然会报错

failed: Not found: Could not find meta graph def matching supplied tags: { serve }. To inspect available tag-sets in the SavedModel, please use the SavedModel CLI: `saved_model_cli`

2. 与saved_model 相关的API

class tf.saved_model.builder.SavedModelBuilder

# 初始化方法
__init__(export_dir)

# 导入graph与变量信息 
add_meta_graph_and_variables(
    sess,
    tags,
    signature_def_map=None,
    assets_collection=None,
    legacy_init_op=None,
    clear_devices=False,
    main_op=None
)

# 载入保存好的模型
tf.saved_model.loader.load(
    sess,
    tags,
    export_dir,
    **saver_kwargs
)

1.1.4 使用SignatureDef：模型载入

载入代码如下：

def test_saved_model_pb(saved_model_path):   
    test_graph = tf.Graph()  
    signature_key = 'test_signature'
    
    with tf.Session(graph=test_graph) as test_sess:
        
        # 调用 tf.saved_model.loader.load 恢复出meta_graph
        meta_graph_def = tf.saved_model.loader.load(
            sess=test_sess, 
            tags=['lstm_saved_model'],  
            export_dir=saved_model_path
        )

        # 从 meta_graph 中取出 Signature 对象
        signature = meta_graph_def.signature_def

        # 从 signature 中取出具体的输入输出的 tensor name
        x_tensor_name = signature[signature_key].inputs['input_XX'].name
        y_tensor_name = signature[signature_key].outputs['pred_YY'].name

        # 取出输入张量和输出张量
        inputXX = test_sess.graph.get_tensor_by_name(x_tensor_name)
        y_pred = test_sess.graph.get_tensor_by_name(y_tensor_name)

        # 加载测试数据
        test_X = ...

        # 得到测试结果
        test_output = test_sess.run([y_pred], feed_dict = {inputXX: test_X})

2. Serving端：模型加载与在线预测

Tensorflow Serving 是google为机器学习模型生产环境部署设计的高性能的服务系统。具有以下特性：

支持模型版本控制和回滚
支持并发与GPU加速，实现高吞吐量
开箱即用，并且可定制化
支持多模型服务
支持 gRPC/ REST API 调用
支持批处理
支持热更新
支持分布式模型
支持多平台模型，如 TensorFlow/MXNet/PyTorch/Caffe2/CNTK等

Tensorflow Serving 丰富的、开箱即用的功能，使得其成为业内认可的部署方案。

Tensorflow Serving 中文官方教程：https://tensorflow.google.cn/tfx/guide/serving
Tensorflow Serving 的Github项目地址：https://github.com/tensorflow/serving

2.1 环境搭建

推荐基于Docker的方式搭建Tensorflow Serving，如何在Ubuntu上安装docker具体可参考我的另外一篇博客：Ubuntu安装及使用Docker

Docker安装完毕后，拉取最新的 tensorflow/serving 的镜像。

docker pull tensorflow/serving

这里直接给出官网（https://tensorflow.google.cn/tfx/serving/docker）示例，运行正常则说明环境搭建完成：

# Download the TensorFlow Serving Docker image and repo
docker pull tensorflow/serving

git clone https://github.com/tensorflow/serving

# Location of demo models
TESTDATA="$(pwd)/serving/tensorflow_serving/servables/tensorflow/testdata"

# Start TensorFlow Serving container and open the REST API port
docker run -t --rm -p 8501:8501 \
    -v "$TESTDATA/saved_model_half_plus_two_cpu:/models/half_plus_two" \
    -e MODEL_NAME=half_plus_two \
    tensorflow/serving &

# Query the model using the predict API
curl -d '{"instances": [1.0, 2.0, 5.0]}' \
    -X POST http://localhost:8501/v1/models/half_plus_two:predict 

# Returns => { "predictions": [2.5, 3.0, 4.5] }

下面对调用预测的API接口的参数进行一个详细解释 :

http://localhost:8501/v1/models/half_plus_two:predict

# 8501: 端口号
# v1: 模型版本号
# half_plus_two : 模型名称
# predict : 模型中的函数名称, 对应于 saved_model 中的 signature_key

        # saved_model 开始，定义 signature_key
        signature_key = 'predict'
        
        # 构建 Signature
        signature = tf.saved_model.signature_def_utils.build_signature_def(...)
        
        # 建立SavedModelBuilder存储模型，并以signature的形式添加要存储的变量
        builder = tf.saved_model.builder.SavedModelBuilder(saved_model_path)
        builder.add_meta_graph_and_variables(
            ... ,
            signature_def_map={signature_key: signature}
            # signature_def_map={'predict':signature}
        )

[补充说明]:

GPU版Tensorflow Serving 的环境搭建见：https://tensorflow.google.cn/tfx/serving/docker

2.2 部署模型

为了部署模型，我们现在要做的是运行Docker容器，将容器的端口发布到主机的端口，然后将主机的路径安装到SavedModel到容器期望模型的位置。

2.2.1 部署单个模型

以上面官的 half_plus_two 模型为例，部署单个模型指令示例：

docker run -p 8501:8501 \
  --mount type=bind,source=/home/xxx/docker/serving/tensorflow_serving/servables/tensorflow/testdata/half_plus_two,target=/models/half_plus_two \
  -e MODEL_NAME=half_plus_two \
  -t tensorflow/serving &

在这种情况下，我们启动了一个Docker容器，将REST API端口8501发布到主机的端口8501，并采用了一个名为 half_plus_two 的模型并将其绑定到默认的模型target路径（ /models/half_plus_two ）。

如果要发布gRPC端口，可以使用-p 8500:8500 。可以同时打开gRPC和REST API端口，或者选择仅打开一个端口。

[补充说明]：

& 表示任务放在后台运行，需要手动kill这个进程。如果不需要后台运行，可以去掉 & 。

# netstat 查看端口占用情况
netstat -pan | grep 8501

# 端口占用情况如下, 占用端口的进程ID为22948
[root@k8s-node1 tensorflow]# netstat -pan |grep 8501
tcp6       0      0 :::8501                 :::*                    LISTEN      22948/docker-proxy

# kill掉ID为22948的进程
kill -9 22948

参数说明：

--mount：   表示要进行挂载
source：    指定要运行部署的模型地址， 也就是挂载的源，这个是在宿主机上的模型目录
target:     这个是要挂载的目标位置，也就是挂载到docker容器中的哪个位置，这是docker容器中的目录
-t:         指定的是挂载到哪个容器
-p:         指定主机到docker容器的端口映射
docker run: 启动这个容器并启动模型服务
 
综合解释：
         将source目录中的例子模型，挂载到-t指定的docker容器中的target目录，并启动

报错分析：
报错1：

invalid argument "type=bind," for "--mount" flag: invalid field '' must be a key=value pair
See 'docker run --help'.

上述shell命令中，逗号后面不可以有空格，可以尝试删除空格，或者去掉换行( \ )。

2.2.2 部署多个模型

通过 config 部署多模型，部署指令示例：

docker run -p 8500:8500 -p 8501:8501 \
  --mount type=bind, \
  source=/tmp/multi_models/, \
  target=/models/multi_models \
  -t tensorflow/serving \ 
  --model_config_file=/models/multi_models/model.config

补充说明：

Serving 镜像支持 gRPC（端口8500）、RESTful API （端口8501）两种方式调用，使用时需要将host的端口与之绑定
Serving 无法直接加载 host 下的模型文件，所以需要将其映射到容器内路径，MODEL_BASE_PATH 默认为 /models
多模型加载和模型版本管理在 model_config_file 中配置

这里给出 model.config 内容示例：

model_config_list:{
  config:{
    name:"textCnn",
    base_path:"/models/multi_models/textCnn/pb",
    model_platform:"tensorflow",
    model_version_policy {
      specific {
        versions: 0
      }
    }
  },

  config:{
    name:"rcnn",
    base_path:"/models/multi_models/rcnn/pb",
    model_platform:"tensorflow",
    model_version_policy {
      specific {
        versions: 0
      }
    }
  },

  config:{
    name:"bert",
    base_path:"/models/multi_models/bert/pb",
    model_platform:"tensorflow",
  }
}