作为一个调参工程师,每日的工作除了清(mo)洗(yu)数(hua)据(shui)外,就是盼望自己呕心沥血堆砌成吨规则的模型能早日上线,让用户早日使用我的“人工”带给ta们“智能”。在NLP领域,一系列预训练模型大放异彩,模型动辄几十上百兆,inference 的时长也是水涨船高。一方面,合理的模型剪枝或蒸馏,可以加速预测速度。另一方面,基于工程的方式提高模型服务的QPS等性能,也是不失为一种策略。
1.前提假设
模型训练时,常规会以 batch 的形式进行。除了训练效果上的考虑,还因为批训练可以节省训练时间,比如某些任务中 batch_size =1 和 batch_size =10 相比,后者的耗时仅为前者的 2-3 倍。预测时同理,如果可以将用户的请求以 batch 形式进行预测,可以大大提升模型服务的效率。
2.实现原理
用户调用模型接口一般都是逐条进行,若将某一时间区间内的若干用户 query 一起组成 batch,送入模型,则可实现”模型一次调用,多个结果输出“的高效率场景。如果可以把接受用户 query 流程和模型预测流程并行化,不互相阻塞,同时以队列作为两流程的通信桥梁,前者作为生产者不断向队列输送 query,后者作为消费者一次获取若干 query 进行预测,即可实现前述场景。
![4fd121f6bdcde3b5ddf49401d05fe749.png](https://img-blog.csdnimg.cn/img_convert/4fd121f6bdcde3b5ddf49401d05fe749.png)
3.实现部分
flask 是一个供 python 用户使用的 web 框架,寥寥几笔就能搭建起一个 demo ,如下
from flask import Flask
app = Flask(__name__)
@app.route('/')
def hello_world():
return 'Hello World'
if __name__ == "__main__":
app.run(host='127.0.0.1', port=8080)
huggingface/transformers 为用户提供了的transformer系列模型的极简调用,同样是寥寥几笔就可以完成情感分类等任务预测,如下
from transformers import pipeline
# Allocate a pipeline for sentiment-analysis
nlp = pipeline('sentiment-analysis')
nlp('We are very happy to include pipeline into the transformers repository.')
>>> {'label': 'POSITIVE', 'score': 0.99893874}
# Allocate a pipeline for question-answering
nlp = pipeline('question-answering')
nlp({
'question': 'What is the name of the repository ?',
'context': 'Pipeline have been included in the huggingface/transformers repository'
})
>>> {'score': 0.28756016668193496, 'start': 35, 'end': 59, 'answer': 'huggingface/transformers'}
Redis 是一款内存数据库,很多开发者也用其作为消息队列。python 可选择的消息队列组件也较多,RabbitMQ、Redis等都可以胜任,这里选择使用Redis作为消息队列。贴上代码
import flask
from transformers import pipeline
import redis
import uuid
import json
from threading import Thread
import time
app = flask.Flask(__name__)
pool = redis.ConnectionPool(host='localhost', port=6379, max_connections=50)
redis_ = redis.Redis(connection_pool=pool, decode_responses=True)
db_key_query = 'query'
db_key_result = 'result'
batch_size = 32
nlp = pipeline("sentiment-analysis")
def classify(batch_size): # 调用模型,设置最大batch_size
while True:
texts = []
query_ids = []
if redis_.llen(db_key_query) == 0: # 若队列中没有元素就继续获取
continue
for i in range(min(redis_.llen(db_key_query), batch_size)):
query = redis_.lpop(db_key_query).decode('UTF-8') # 获取query的text
query_ids.append(json.loads(query)['id'])
texts.append(json.loads(query)['text']) # 拼接若干text 为batch
result = nlp(texts) # 调用模型
for (id_, res) in zip(query_ids, result):
res['score'] = str(res['score'])
redis_.set(id_, json.dumps(res)) # 将模型结果送回队列
@app.route("/predict", methods=["POST"])
def handle_query():
text = flask.request.form['text'] # 获取用户query中的文本 例如"I love you"
id_ = str(uuid.uuid1()) # 为query生成唯一标识
d = {'id': id_, 'text': text} # 绑定文本和query id
redis_.rpush(db_key_query, json.dumps(d)) # 加入redis
while True:
result = redis_.get(id_) # 获取该query的模型结果
if result is not None:
redis_.delete(id_)
result_text = {'code': "200", 'data': result.decode('UTF-8')}
break
return flask.jsonify(result_text) # 返回结果
if __name__ == "__main__":
t = Thread(target=classify, args=(batch_size,))
t.start()
app.run(debug=False, host='127.0.0.1', port=9000)
整个流程分为两个子流程,其一是处理用户 query 流程,其二是调用模型产生结果流程。
handle_query() 方法首先获取 request 的 text,为该 text 配置 id, 并推入队列。随后便等待模型返回结果,获得结果后,将该 text 与 id 从队列中删除,而后为用户返回结果。
classify() 方法调用模型预测结果,首先从队列中获取最多 batch_size 个 text ,拼接成为 batch 随后调用模型,模型给出结果后,再次放入队列中,供 handle_query() 获取结果并返回用户。
4.速度对比
分别用10、15、20、30 线程并发请求模型服务,每个线程发送 query 50条,耗时结果如下,蓝色为没有加入 Redis 的耗时情况,耗时增幅很快。橙色为加入Redis,批量预测耗时,基本在60秒左右震荡,增幅较小。
![2079468c0abd1b28b0e6947696a0e9e0.png](https://img-blog.csdnimg.cn/img_convert/2079468c0abd1b28b0e6947696a0e9e0.png)
5.结论
本文初步探索了 Redis 做为消息队列,利用深度模型批次预测的特点,提升在线模型服务的性能。希望各位大佬,提出宝贵意见,共同探讨。
【参考资料】:
- Redis【入门】就这一篇!
- A scalable Keras + deep learning REST API
- hugging face
- 初识flask,搭建第一个自己的网页