pytorch - 简单的 MultiGPU 在推理过程中使用 huggingface

Question

我有两个 GPU。

如何将它们用于拥抱脸管道的推理？

Huggingface 文档似乎说我们可以轻松地将 DataParallel 类与 huggingface 模型一起使用，但我没有看到任何示例。

例如使用 pytorch，只需执行以下操作就非常容易:

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
output = net(input_var)  # input_var can be on any device, including CPU 

是否有与 huggingface 等价的东西？

Answer 1

我发现管道不可能，所以:

两种方式:

• 用 huggingface 中的 Trainer 对象来做，它也支持推理，但不是最优的。
• 使用多处理标准库中的队列，但这会产生大量样板代码