深度学习模型通常实现为通过将输入组合成“小批量”来有效地使用 GPU。然而，为这些模型提供服务的应用程序通常逐个接收请求。因此，使用传统的单线程或多线程服务器方法会导致 GPU 资源利用率低下，并且会导致随着请求数量的增加而线性增加的延迟。

batched-fn 是一个针对深度学习 web 服务器的即插即用解决方案，它将单个请求排队，并将它们作为批量提供给您的模型。只需将 batched_fn 宏插入到运行模型请求的函数中，即可将其添加到任何应用程序中，而无需进行最小程度的重构。

特性

• 🚀 易于使用：将 batched_fn! 宏直接放入现有代码中。
• 🔥 轻量级且快速：在闪电般的快速 flume crate 之上实现了队列系统。
• 🙌 易于调整：只需调整 max_delay 和 max_batch_size。
• 🛑 包含 背压 机制：只需设置 channel_cap 并通过从您的 web 服务器返回 503 来处理 Error::Full。

示例

假设您有一个模型 API，如下所示

// `Batch` could be anything that implements the `batched_fn::Batch` trait.
type Batch<T> = Vec<T>;

#[derive(Debug)]
struct Input {
    // ...
}

#[derive(Debug)]
struct Output {
    // ...
}

struct Model {
    // ...
}

impl Model {
    fn predict(&self, batch: Batch<Input>) -> Batch<Output> {
        // ...
    }

    fn load() -> Self {
        // ...
    }
}

在没有 batched-fn 的情况下，web 服务器路由需要对每个单个输入调用 Model::predict，从而导致由于 GPU 利用率低下而形成的瓶颈

use once_cell::sync::Lazy;
static MODEL: Lazy<Model> = Lazy::new(Model::load);

fn predict_for_http_request(input: Input) -> Output {
    let mut batched_input = Batch::with_capacity(1);
    batched_input.push(input);
    MODEL.predict(batched_input).pop().unwrap()
}

但是，通过将batched_fn 宏添加到您的代码中，您可以在不改变输入和输出一对一关系的情况下，在幕后自动获得批量推理。

async fn predict_for_http_request(input: Input) -> Output {
    let batch_predict = batched_fn! {
        handler = |batch: Batch<Input>, model: &Model| -> Batch<Output> {
            model.predict(batch)
        };
        config = {
            max_batch_size: 16,
            max_delay: 50,
        };
        context = {
            model: Model::load(),
        };
    };
    batch_predict(input).await.unwrap()
}

❗️ 请注意，现在 predict_for_http_request 函数必须是 async。

在这里，我们将 max_batch_size 设置为 16，并将 max_delay 设置为 50 毫秒。这意味着批量函数将在接收单个输入后最多等待 50 毫秒来填充一个 16 个元素的批量。如果在 50 毫秒内没有接收到 15 个更多的输入，则将按原样运行部分批量。

调整最大批量大小和最大延迟

最佳批量大小和延迟将取决于您的特定用例，例如您可以在内存中容纳多大（对于深度学习模型，通常是 8、16、32 或 64）的批量以及您能够承受多长的延迟。一般来说，您希望将 max_batch_size 设置得尽可能高，假设以 N 的大小进行批量处理时，N 个示例的总处理时间最小化，同时将 max_delay 与处理一个批量的时间相比保持较小。

实现细节

当调用 batched_fn 宏时，它会启动一个新线程，在该线程中运行 handler。在该线程中，初始化 context 中指定的每个对象，然后每次运行时通过引用传递给处理程序。

宏返回的对象只是一个闭包，它通过异步通道将单个输入和回调发送到处理程序线程。当处理程序完成运行一个批量的操作后，它会调用与每个输入对应的回调，并传递相应的输出，从而触发闭包唤醒并返回输出。