基于TensorFlow Serving的YOLO模型部署

在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服务部署和调用

由于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倍左右。

哈喽，喜欢这篇文章的话烦请点个赞哦！万分感谢_{(^▽^)PS：有问题可以联系我们哦}v ceshiren001

复制“下方链接”，提升测试核心竞争力！
更多技术文章分享和免费资料领取