SBERT之部署TF-serving和Flask服务实践详解

背景

SBERT模型是基于pytorch和transformers来实现的。如果想要通过tensorflow-serving来部署模型的话。首先，要将transformers模型转成pb的。tensorflow-serving只能支持模型本身的输入输出，不支持模型前的预处理，模型后的pooling等业务相关定制化操作。基于我们的定制化业务需求，决定了tensorflow-serving上需要封装一层flask服务来实现定制化业务。从而实现对客户端而言是端到端的输入与输出，将业务封装在服务端。

tensorflow serving部署

使用docker部署tensorflow serving参考：TensorFlow Serving with Docker

启动tensorflow serving的容器：

docker run -t --rm -p 8501:8501 \
    -v "$TESTDATA/saved_model_half_plus_two_cpu:/models/half_plus_two" \
    -e MODEL_NAME=half_plus_two \
    tensorflow/serving &

验证tensorflow serving功能：

通过curl验证

curl -d '{"instances": [1.0, 2.0, 5.0]}' \
    -X POST http://localhost:8501/v1/models/half_plus_two:predict

通过postman验证

启动优化后的SBERT模型tensorflow serving容器：

docker run -t --rm -p 8501:8501 -v "/home/models/paraphrase-multilingual-mpnet-base-v2.pb:/models/paraphrase-multilingual-mpnet-base-v2" -e MODEL_NAME=paraphrase-multilingual-mpnet-base-v2 tensorflow/serving &

踩坑记录

1

—— 部署到tensorflow serving的模型必须要有子文件夹标注版本号

模型名/ 版本号（如：1、2、3等数字）/ saved_model.pb

2

在调试tensorflow serving输入数据时使用instances时，报错：

The second input must be a scalar, but it has shape [14]\n\t [[{{node assert_equal_5/All}}]]

https://www.tensorflow.org/tfx/serving/api_rest

—— 按照说明里instances应该可以用，可实际上报错，换成inputs可以用

参考：Tensorflow Serving-Docker RESTful API客户端访问问题排查 - 简书

flask服务部署

制作flask镜像

参考：用Docker部署Flask应用

mkdir flask

cd flask

创建如下文件：Dockerfile、gunicorn.conf.py、requirements.txt、业务相关文件

Dockerfile

FROM python:3.9 WORKDIR /usr/src/app COPY requirements.txt ./ RUN pip install -i https://pypi.tuna.tsinghua.edu.cn/simple --upgrade pip RUN pip install -i https://pypi.tuna.tsinghua.edu.cn/simple --no-cache-dir -r requirements.txt COPY . . CMD ["gunicorn", "app:app", "-c", "./gunicorn.conf.py"]

gunicorn.conf.py

# -*- coding: UTF-8 -*- workers = 5    # 定义同时开启的处理请求的进程数量，根据网站流量适当调整 worker_class = "gevent"   # 采用gevent库，支持异步处理请求，提高吞吐量 bind = "0.0.0.0:8888"    # 监听IP放宽，以便于Docker之间、Docker和宿主机之间的通信

requirements.txt

flask gunicorn gevent requests transformers pyyaml==5.4.1 tensorflow annoy pickle-mixin torch torchvision torchaudio

执行镜像创建命令：

docker build -t 'flask' .

启动flask服务

创建flask服务容器

docker run -it --rm -p 8888:8888 -v "/home/wyf/tokenizer:/tokenizer" -v "/home/wyf/index:/index" flask &

业务的实现是在app.py中封装实现的

• SBERT模型输入预处理（利用transformers自己的tokenizer，底层需要安装pytorch支持）
• pb模型处理（tensorflow-serving部署）
• 通过自定义pooling加工出最后的embeddings

踩坑记录

1

pip install时指定国内镜像源加如下参数

-i  https://pypi.mirrors.ustc.edu.cn/simple/

2

error pulling image configuration dial tcp i/o timeout

—— docker拿不到image，网络问题，关闭防火墙

service ufw status

3

app.run(debug=True, host="0.0.0.0")

一定要放在main里面，否则启动flask时报错

4

ValueError: unsupported pickle protocol: 5

——保存pickle的python版本和读取pickle的python版本要一致（在我的情况下高版本python读取低版本python保存的pickle也可以，貌似向下兼容，反之报错）

5

带中文的python文件要是运行时报解码错误

—— 在文件开始加入

# -*- coding: UTF-8 -*-

相关代码参考：GitHub - tdaajames/aitest