基于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倍左右。
参考文献
https://www.tensorflow.org/tfx/guide/serving
https://www.tensorflow.org/guide/saved_model
长按二维码，关注官方微信公众号！
360开测
QQ交流群
保存二维码
再用QQ扫描哟
温馨提示
如果你喜欢本文，请分享到朋友圈，想要获得更多信息，请关注我。