基于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

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

900×382 111 KB

安装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