借助批量推理提升推理效率
2020 年冬
作为惠普 Z 系列全球数据科学大使，Yuanhao Wu 的内容由惠普赞助，且向他提供了惠普产品。
作为一名算法工程师，我在日常工作中经常需要完成模型上线任务。对于一些要求较低的场景，可以利用网络框架完成上线任务：针对每个用户请求，调用模型进行推理并返回结果。然而，这种简单的实现通常不能最大化 GPU 的使用效率，并且无法应付拥有高性能需求的场景。
我们可以通过许多方法来进行优化，其中一个很有用的技巧便是将针对每个请求进行推理改为一次针对多个请求进行推理。大约在去年的这个时候，我编写了一个小工具来实现这个功能，还给它起了一个相当霸气的名字：InferLight。老实说，这个工具的实现效果并不理想。最近，我参照 Shannon Technology 的 Service-Streamer 对其进行了重构。
这个工具的功能看似简单，但在部署的过程中，我们可以了解很多 Python 异步编程的知识，感受现代 GPU 的并行计算能力。
架构
首先，为了提高模型的在线推理通量，我们需要使推理服务异步化。对于网络服务来说，异步意味着程序可以在模型执行计算时处理其他请求。对于 Python，我们可以通过基于 Asyncio 的出色框架实现异步服务，比如我经常使用的 Sanic。虽然推理是一项计算密集型任务，但我们的目标是能够聚合多个推理请求，有效地利用 GPU 的并行计算能力，同时正确地将批量推理的结果返回给相应的请求者。
为了实现上述目标，需要以下模块：
- 前端服务：用于接收请求和返回结果。我们可以使用各种协议，如 Http、PRC 等，这是一个独立的进程。
- 推理工作器：负责模型初始化、批量推理数据构建和推理计算。这也是一个独立的进程。
- 任务队列：前端服务接收请求，然后将计算任务发送到任务队列；推理工作器侦听队列，每次取出一小批任务交由模型进行推理
- 结果队列：推理完成后，推理工作器将结果发送到结果队列；前端服务侦听结果队列并获取推理结果
- 结果分发：在将任务发送到任务队列前，需要生成任务的唯一标识符，以便从结果队列取回结果后根据该标识符获取与任务对应的结果
我们有很多方法来实现任务队列和结果队列，可以使用一些成熟的中间件，比如 Kafka 和 Redis。为了避免外部依赖，这次我选择使用 Python 的原生多进程队列。结果队列侦听和分发通过前端服务进程的一个子线程来完成。
部署
推理工作器相对简单。由于要加载各种各样的模型且数据处理步骤种类繁多，所以我将推理工作器设计成一个继承而来的基类，在使用时需要实现特定方法。
与之相搭配的是前端服务中使用的包装类，用于处理推理请求的接收、结果收集和分发。
其中用到的一些数据结构定义如下：
使用案例和测试结果
此处展示如何将上述组件与情感分析 BERT 模型配合使用。
首先，定义模型
接着，继承 BaseInferLightWorker 并实现三个函数来获得一个完整的工作器分类
最后，构建服务
我用知名的 Apache ab 工具做了一些测试。我在 HP Z4 工作站上启动了上述应用，并保证确保工作器进程是在 RTX 6000 GPU 上运行的。
通过使用 ab -n 1000 -c 32 http://localhost:8888/batched_predict，我得到了以下结果。
另一个简单的实现并未借助批量推理，其测试结果如下所示：
如大家所见，通过批量推理，我们大约将通量提升到了原来的 2.5 倍！在进行基准测试时，我还观察到使用批量推理可以大幅提升 GPU 使用率。
我已经公布了 InferLight 的源代码，可点击此处查看。希望大家喜欢 :)
有疑问？
联系销售支持。
不知道该怎么选
Z 系列工作站？
需要 Z 系列工作站方面的支持？
免责声明
产品的Logo及位置可能与图片有所差异，不影响产品性能和功能，请以实物为准。
产品图片仅供参考，因不同国家可能存在稍许差异，实际产品以销售为准。
本文所载信息如有变更，恕不另行通知。惠普产品与服务的完整保修条款见此类产品和服务附带的正式保修声明。本文中的信息概不构成额外的保修条款。惠普对本文包含的技术或编辑方面的错误或遗漏概不负责。