PyTorch:如何批量推理(并行推理)

Question

如何在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

目前正在尝试理解该库并创建一个工作示例。在此期间，您是否知道更好的方法？

Answer 1

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

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