TensorFlow Serving 模型部署与实践

前言

在《基于 TensorFlow Serving 的 YOLO 模型部署》文章中有介绍 tensorflow 1.x 版本的模型如何利用 TensorFlow Serving 部署。本文接着上篇介绍 tensorflow2.x 版本的模型部署。

工作原理

架构图

核心概念

servables:对机器学习模型的封装和抽象，可以提供模型计算服务，一般有 HTTP restful server + grpc server + 模型管理组成。一个机器学习模型可以表示成一个或多个 servable
sources：从指定文件目录中发现模型，封装成 servable,每个 source 提供一个或者多个 servable 流（一系列特定版本的 servable 序列），并为每个版本提供一个 loader 实例，使得其变成可加载的
loader: 提供加载卸载 servable 的 API
Aspired version: 代表需要被加载或者已经就绪的 servable 集合。由 source 提供，manager 负责管理该集合。
Manager: 负责管理 servable 的整个生命周期，manager 监听 sources 并跟踪所有版本的 servable,manager 试图满足所有 source 的需求，但如果发现系统资源无法满足待加载的 servable 时，也会拒绝加载。Manager 根据 version policy 来管理 servable
VersionPolicy：tfserving 有两个默认版本控制策略，1) Availability Presrving Policy—避免没有 version 加载的情况，即在卸载一个旧 version 前通常会加载一个新的 version; 2) Resource Preserving Policy—避免两个 version 同时被加载，这需要双倍的资源
ServableHandler: servable 实例，用于处理 client 发送的请求

servable 的生命周期：

● 一个 Source 插件会为一个特定的 version 创建一个 Loader。该 Loaders 包含了在加载 Servable 时所需的任何元数据。
● Source 会使用一个 callback 来通知该 Aspired Version 的 Manager
● 该 Manager 应用该配置过的 Version Policy 来决定下一个要采用的动作，它会被 unload 一个之前已经加载过的 version，或者加载一个新的 version
● 如果该 Manager 决定它是否安全，它会给 Loader 所需的资源，并告诉该 Loader 来加载新的 version
● Clients 会告知 Manager，即可显式指定 version，或者只请求最新的 version。该 Manager 会为该 Servable 返回一个 handle

Tensorflow 2.x 模型部署

TF serving 环境准备
Tensorflow Serving 环境最简单的安装方式是 docker 镜像安装。
拉取 CPU 版的 tensorflow/serving 镜像：

docker pull tensorflow/sering:last

拉取 GPU 版的 tensorflow/serving 镜像：
安装 GPU 版 docker：nvidia-docker，nvidia-docker 是在 docker 上做一层封装，通过 nvidia-docker-plugin,调用到 docker 上，ubuntu 上安装命令

sudo apt-get install -y nvidia-docker2

从 docker 仓库中获取 GPU 版本的镜像:

docker pull tensorflow/serving:latest-devel-gpu

模型保存—savedmodel
Tensorflow 2.x 模型有以下几种格式：
1.checkpoint 格式

model.save_weights(“./xxx.ckpt” , save_format=”tf”)

2.h5 格式

model.save(“./xxx.h5”)
model.save_weights(“./xxx.h5”, save_format=”h5”)

3.saved_model 格式

model.save(“./xxx”, save_format=”tf”)
tf.saved_model.save(obj, “./xxx”)

4.保存模型结构

model.to_json()

tf serving 支持 saved_model 格式的模型部署，saved_model 格式模型示意图如下图所示，.pb 存储了模型的结构信息和常量，variables 存储了模型变量数据，assets 用于存储额外的初始化所需文件，如词典。

例子：

# coding=utf-8
import tensorflow as tf

class TestTFServing(tf.Module):
    def __init__(self):
        self.x = tf.Variable("hello", dtype = tf.string,trainable=True)

    @tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.string)])
    def concat_str(self, a):
        self.x = self.x + a
        return self.x

    @tf.function(input_signature=[tf.TensorSpec(shape = [], dtype = tf.string)])
    def cp_str(self, b):
        self.x.assign(b)
        return self.x


if __name__ == '__main__':
    demo = TestTFServing()
    tf.saved_model.save(demo, "model\\test\\1", signatures={"test_assign": demo.cp_str,\
     "test_concat": demo.concat_str})

# coding=utf-8

import tensorflow as tf

class DenseNet(tf.keras.Model):
    def __init__(self):
        super(DenseNet, self).__init__()

    def build(self, input_shape):
        self.dense1 = tf.keras.layers.Dense(15, activation='relu')
        self.dense2 = tf.keras.layers.Dense(10, activation='relu')
        self.dense3 = tf.keras.layers.Dense(1, activation='sigmoid')
        super(DenseNet, self).build(input_shape=input_shape)

    def call(self, x):
        x = self.dense1(x)
        x = self.dense2(x)
        x = self.dense3(x)
        return x


if __name__ == '__main__':
    model = DenseNet()
    model.build(input_shape=(None, 15))
    model.summary()
    inputs = tf.random.uniform(shape=(10, 15))
    model._set_inputs(inputs=inputs) # tf2.0 need add this line
    model.save(".\\model\\test\\2", save_format="tf")
    tf.keras.models.save_model(model, ".\\model\\test\\3", save_format="tf")

服务启动

CPU 版模型部署命令：

docker run -p 8500:8500 -p 8501:8501 --mount "type=bind,source=/home/test/ybq/model/demo,target=/models/demo" -e MODEL_NAME=demo tensorflow/serving:latest

GPU 版模型部署命令：

docker run -p 8500:8500 -p 8501:8501 --runtime nvidia --mount "type=bind,source=/home/test/ybq/model/demo,target=/models/demo" -e MODEL_NAME=demo tensorflow/serving:latest-gpu

Warm up 模型

由于 tensorflow 有些组件是懒加载模式，因此第一次请求预测会有很严重的延迟，为了降低懒加载的影响，需要在服务初始启动的时候给一些小的请求样本，为了降低懒加载的影响，需要在服务初始启动的时候给一些小的样本 (tfrecord 格式)，调用模型的预测接口，预热模型。
WarmUp Model 步骤如下所示：
1.生成 Warmup 数据，文件 tf_serving_warmup_requests，格式 tfrecord

# coding=utf-8
import tensorflow as tf
from tensorflow_serving.apis import model_pb2
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_log_pb2
def main():
    with tf.io.TFRecordWriter("tf_serving_warmup_requests") as writer:
        request = predict_pb2.PredictRequest(
            model_spec=model_pb2.ModelSpec(name="demo", signature_name='serving_default'),
            inputs={"x": tf.make_tensor_proto(["warm"]), 
                    "y": tf.make_tensor_proto(["up"])}
        )
        log = prediction_log_pb2.PredictionLog(
            predict_log=prediction_log_pb2.PredictLog(request=request))
        writer.write(log.SerializeToString())
if __name__ == "__main__":
    main()

2.将生成的 tfrecord 放在对应模型和版本所在目录 assets.extra 文件夹下

模型维护

默认版本维护策略是只会加载同时加载 servable 的一个版本，但是我们可以通过配置 config，修改模型版本加载策略，也可以通过配置多个模型，同时维护多个模型。
模型版本配置：

model_config_list {
  config {
    name: "demo"
    base_path: "/models/demo"
    model_platform: "tensorflow"
    model_version_policy{
        specific {
            versions: 1
            versions: 2
        }
    }
  }
}

将 model.config 放到 model 目录下，启动命令如下：

docker run -p 8500:8500 -p 8501:8501 --mount "type=bind,source=/home/test/ybq/model/demo,target=/models/demo" -e MODEL_NAME=demo \
tensorflow/serving:latest \
--model_config_file=/models/demo/model.config \
--model_config_file_poll_wait_seconds=60

多个模型部署配置：

model_config_list {
  config {
    name: "demo"
    base_path: "/models/model/demo"
    model_platform: "tensorflow"
    model_version_policy{
        specific {
            versions: 1
            versions: 2
        }
    }
  }
  config {
    name: "111"
    base_path: "/models/model/111"
    model_platform: "tensorflow"
    model_version_policy{
        specific {
            versions: 1
        }
    }
  }
}

将 model.configs 放到 model 目录下，启动命令如下：

docker run  -p 8500:8500 -p 8501:8501 --mount "type=bind,source=/home/test/ybq/model/,target=/models/model" tensorflow/serving:latest \
--model_config_file=/models/model/models.config \
--model_config_file_poll_wait_seconds=60

服务调用

TF Serving 提供了两种方式，一种是 grpc 方式（默认端口是 8500），另一种是 http 接口调用方式（默认端口是 8501）。
对于 tensorflow2.x 版本生成的 saved_model 模型，没有像 1.x 版本使用 SavedModelBuilder API 自定义签名（输入输出的数据类型 + 方法），2.x 版本模型输入的数据类型可以通过@tf.function中的 input_signature 参数指定，方法目前来看是写死在源码中的，只有 signature_constants.PREDICT_METHOD_NAME 一种。很多时候模型的输入输出对我们来说是黑盒的，而在调用服务接口的时候我们需要知道模型的输入输出以及签名的 key，我们可以使用 saved_model_cli show --dir model/test/1 –all 来查看我们需要的参数

grpc 和 http 请求对应的代码片段分别如下所示：

# coding=utf-8
import requests
import json
import tensorflow as tf
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
import grpc

def test_grpc():
    channel = grpc.insecure_channel('{host}:{port}'.format(host="127.0.0.1", port=8500))
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    request = predict_pb2.PredictRequest()
    request.model_spec.name = "demo"
    request.model_spec.signature_name = "test_concat"
    request.inputs['a'].CopyFrom(tf.make_tensor_proto("xxx"))
    result = stub.Predict(request, 10.0)
    return result.outputs


def test_http():
    params = json.dumps({"signature_name": "test_concat", "inputs": {"a": "xxx"}})
    data = json.dumps(params)
    rep = requests.post("http://127.0.0.1:8501/v1/models/demo/version1/:predict", data=data)
    return rep.text

参考
Tensorflow