8在线 form girls_使用Redis加速深度学习模型的在线预测

作为一个调参工程师，每日的工作除了清(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 进行预测，即可实现前述场景。

消息队列，摘自“基于Redis实现消息队列的典型方案”

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秒左右震荡，增幅较小。

耗时对比

5.结论

本文初步探索了 Redis 做为消息队列，利用深度模型批次预测的特点，提升在线模型服务的性能。希望各位大佬，提出宝贵意见，共同探讨。

【参考资料】：

1. Redis【入门】就这一篇！
2. A scalable Keras + deep learning REST API
3. hugging face
4. 初识flask，搭建第一个自己的网页