批量推理工具包

批量推理工具包(batch-inference)是一个Python包，它可以动态地将来自多个用户的模型输入张量进行批处理，执行模型，然后分别将输出张量解批处理并返回给每个用户。这将通过更好的计算并行性和更好的缓存局部性来提高系统吞吐量。整个过程对开发者来说是透明的。

何时使用

当您想在云服务器上托管深度学习模型推理时，尤其是在GPU上时

为什么要使用

它可以提高您的服务器吞吐量多达几倍

批量推理的优势

• 平台独立的轻量级Python库
• 只需少量代码更改即可使用内置的批处理算法加入
• 灵活的API，支持自定义批处理算法和输入类型
• 支持多进程远程模式，以避免Python GIL瓶颈
• 关于流行模型的教程和基准测试

安装

从Pip安装

python -m pip install batch-inference --upgrade

从源码构建和安装 (适用于开发者)

git clone https://github.com/microsoft/batch-inference.git
python -m pip install -e .[docs,testing]

# if you want to format the code before commit
pip install pre-commit
pre-commit install

# run unittests
python -m unittest discover tests

示例

让我们从一个玩具模型开始，学习API。首先，您需要在您的模型类中定义一个 predict_batch 方法，然后添加 batching 装饰器到您的模型类。

批量装饰器为创建 ModelHost 对象添加了 host() 方法。ModelHost 的 predict 方法接收单个查询作为输入，并在调用 predict_batch 方法之前将多个查询合并成一个批次。predict 方法还会在返回结果之前将 predict_batch 方法的输出拆分。

import numpy as np
from batch_inference import batching
from batch_inference.batcher.concat_batcher import ConcatBatcher

@batching(batcher=ConcatBatcher(), max_batch_size=32)
class MyModel:
    def __init__(self, k, n):
        self.weights = np.random.randn((k, n)).astype("f")

    # shape of x: [batch_size, m, k]
    def predict_batch(self, x):
        y = np.matmul(x, self.weights)
        return y

# initialize MyModel with k=3 and n=3
host = MyModel.host(3, 3)
host.start()

# shape of x: [1, 3, 3]
def process_request(x):
    y = host.predict(x)
    return y

Batcher 负责合并查询和拆分输出。在这种情况下，ConcatBatcher 会将输入张量在第一个维度上拼接成批处理张量。我们提供了一套内置的 Batcher 以满足常见场景，您也可以实现自己的 Batcher。有关更多信息，请参阅 什么是 Batcher。