背景

在 UI 自动化测试中，界面控件识别是基石。在计算机视觉领域中，有很多经典的目标识别模型，我们尝试将 YOLO 模型迁移至自动化测试领域界面控件识别中。迁移训练后的模型需要部署到生产环境，TensorFlow Serving 是一种模型部署方法，只需几行简单的代码就可以维护模型的整个生命周期。下面将以原 YOLO V3 tensorflow 版模型为例从环境准备、模型格式转换、服务部署和调用以及最后的性能对比四个方面介绍 TensorFlow Serving 在模型部署中的应用。

环境准备

官网推荐 TensorFlow Serving 在 docker 容器中运行，因此安装 TensorFlow Serving 之前需要先安装 docker。
安装 TensorFlow Serving cpu 版本步骤：
从 docker 仓库中获取 cpu 版本的镜像：

docker pull tensorflow/serving

启动容器：

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

启动成功示意图如下图所示：

安装 Tensorflow Serving GPU 版本步骤：
安装 nvidia-docker

sudo apt-get install -y nvidia-docker2 

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

docker pull tensorflow/serving:latest-devel-gpu

启动容器：

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

启动成功的示意图如下图所示：

YOLO 模型格式转换

YOLOV3 tensorflow 版本保存的模型格式是 ckpt 格式，tensorflow serving 需要的格式是 savedmodel 格式，因此需要进行模型格式转换。ckpt 格式的模型只存储一些权重，我们需要创建一个 session，将计算图搭建好，并恢复 ckpt 中保存的参数，定义标记模型的输入和输出参数签名，从而得到 savedmodel 格式的模型，具体实现代码如下图所示：

with tf.Session() as sess:
    string_inp = tf.placeholder(tf.string, shape=(None,))  # string input for the base64 encoded image
    imgs_map = tf.map_fn(
        tf.image.decode_image,
        string_inp,
        dtype=tf.uint8
    )  # decode jpeg
    imgs_map.set_shape((None, None, None, 3))
    imgs = tf.image.resize_images(imgs_map, [416, 416])  # resize images
    imgs = tf.reshape(imgs, (-1, 416, 416, 3))  # reshape them
    img_float = tf.cast(imgs, dtype=tf.float32) / 255
    yolo_model = yolov3(num_class, anchors)
    with tf.variable_scope('yolov3'):
        pred_feature_maps = yolo_model.forward(img_float, False)
    pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)

    pred_scores = pred_confs * pred_probs

    boxes, scores, labels = gpu_nms(pred_boxes, pred_scores, num_class, max_boxes=200, score_thresh=0.3, nms_thresh=0.45)
    saver = tf.train.Saver()
    saver.restore(sess, "./data/darknet_weights/yolov3.ckpt")
    print('Exporting trained model to', export_path)
    # 构造定义一个builder，并制定模型输出路径
    builder = tf.saved_model.builder.SavedModelBuilder(export_path)
    # 声明模型的input和output
    tensor_info_input = tf.saved_model.utils.build_tensor_info(string_inp)
    tensor_info_output1 = tf.saved_model.utils.build_tensor_info(boxes)
    tensor_info_output2 = tf.saved_model.utils.build_tensor_info(scores)
    tensor_info_output3 = tf.saved_model.utils.build_tensor_info(labels)
    # 定义签名
    prediction_signature = (
        tf.saved_model.signature_def_utils.build_signature_def(
            inputs={'images': tensor_info_input},
            outputs={"boxes": tensor_info_output1, "scores": tensor_info_output2, "labels": tensor_info_output3},
            method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))
    builder.add_meta_graph_and_variables(sess, [tf.saved_model.tag_constants.SERVING],
                                         signature_def_map={'predict_images': prediction_signature,
                                                            tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                                                                prediction_signature})
    builder.save()
    print('Done exporting!')

YOLO 服务部署和调用

Warmup Model：
由于 tensorflow 有些组件是懒加载模式，因此第一次请求预测会有很严重的延迟，为了降低懒加载的影响，需要在服务初始启动的时候给一些小的请求样本，调用模型的预测接口，预热模型。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():
    data = (open('./data/dog.jpg', 'rb').read())
    with tf.io.TFRecordWriter("tf_serving_warmup_requests") as writer:
        request = predict_pb2.PredictRequest(
            model_spec=model_pb2.ModelSpec(name="yolo", signature_name='predict_images'),
            inputs={"images": tf.make_tensor_proto([data])}
        )
        log = prediction_log_pb2.PredictionLog(
            predict_log=prediction_log_pb2.PredictLog(request=request))
        writer.write(log.SerializeToString())


if __name__ == "__main__":
    main()

2、将生成的 tf_serving_warmup_requests 文件放到对应模型和版本所在目录 assets.extra 文件夹下
3、启动模型时，模型会自动加载 tf_serving_warmup_requests 中的数据预热模型，预热模型的日志输出示意图如下图所示：

版本维护：
tf serving 可以维护模型的多个版本，当在模型目录下放入新版本模型，tf serving 监听到有新版本模型加入后，会将新版本模型加载进内存中同时卸载之前加载的旧模型，日志输出示意图如下图所示：

服务调用：
调用部署好的 yolo 模型，tf serving 提供了两种方式，一种是 grpc 方式（默认端口是 8500），另一种是 http 接口调用方式（默认端口是 8501）。grpc 和 http 请求对应的代码片段分别如下图所示：

def test_grpc(img_path):
    channel = grpc.insecure_channel('{host}:{port}'.format(host="10.18.131.58", port=8500))
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    data = (open(img_path, 'rb').read())
    request = predict_pb2.PredictRequest()
    request.model_spec.name = "yolo"
    request.model_spec.signature_name = "predict_images"
    request.inputs['images'].CopyFrom(tf.make_tensor_proto([data]))
    result = stub.Predict(request, 10.0)
    return result


def test_http(img_path):
    with open(img_path, "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read())
    params = {"inputs": {"images": [{"b64": encoded_string}]}}
    data = json.dumps(params)
    rep = requests.post("http://10.18.131.58:8501/v1/models/yolo:predict", data=data)
    return rep.text

CPU&GPU 性能对比
我们对比了 TensorFlow Serving 仅使用 CPU 和使用 GPU 加速两种情况下 YOLO 服务的性能，对比发现，使用 GPU 加速后处理一次请求的时间是 50ms 左右，不使用 GPU 加速处理一起请求需要 280ms 左右。响应速度提升 5~6 倍左右。

参考文献
【1】https://www.tensorflow.org/tfx/guide/serving
【2】https://www.tensorflow.org/guide/saved_model