PyTorch：如何批量进行推理（并行推理）

Posted 2023-02-16

【中文标题】PyTorch：如何批量进行推理（并行推理）

【英文标题】：PyTorch: How to do inference in batches (inference in parallel)

【发布时间】：2020-08-26 18:31:02

【问题描述】：

如何在 PyTorch 中进行批量推理？如何并行进行推理以加快这部分代码的速度。

我已经开始使用标准的推理方式：

with torch.no_grad():
    for inputs, labels in dataloader['predict']:
        inputs = inputs.to(device)
        output = model(inputs)
        output = output.to(device)

我研究过，唯一提到并行推理（在同一台机器上）似乎是使用 Dask 库：https://examples.dask.org/machine-learning/torch-prediction.html

目前正在尝试了解该库并创建一个工作示例。同时你知道更好的方法吗？

【问题讨论】：

如果inputs 是一批输入，那么您的代码会分批进行推理。如果您认为它很慢，请首先检查 GPU 是否被充分利用（内存和利用率）。让我们知道。

【参考方案1】：

在 pytorch 中，输入张量总是在第一个维度中具有批处理维度。因此按批次进行推理是默认行为，您只需将批次维度增加到大于 1。

例如，如果您的单个输入是[1, 1]，则其输入张量为[[1, 1], ]，形状为(1, 2)。如果您有两个输入 [1, 1] 和 [2, 2]，则将输入张量生成为 [[1, 1], [2, 2], ]，形状为 (2, 2)。这通常在批处理生成器函数中完成，例如您的dataloader。