背景

在 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 imageimgs_map = tf.map_fn(tf.image.decode_image,string_inp,dtype=tf.uint8) # decode jpegimgs_map.set_shape((None, None, None, 3))imgs = tf.image.resize_images(imgs_map, [416, 416]) # resize imagesimgs = tf.reshape(imgs, (-1, 416, 416, 3)) # reshape themimg_float = tf.cast(imgs, dtype=tf.float32) / 255yolo_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和outputtensor_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 的步骤如下所示：

Warmup 数据的生成代码如下图所示，运行后生成

tf_serving_warmup_requests.TFRecord。

coding=utf-8

import tensorflow as tffrom tensorflow_serving.apis import model_pb2from tensorflow_serving.apis import predict_pb2from 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()

将生成的 tf_serving_warmup_requests 文件放到对应模型和版本所在目录 assets.extra 文件夹下启动模型时，模型会自动加载 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 倍左右

点击下方链接免费领取：性能测试+接口测试+自动化测试+测试开发+测试用例+简历模板+测试文档

http://qrcode.testing-studio.com/f?from=infoQ&url=https://ceshiren.com/t/topic/22265