如何从 pytorch 模型并行化模型预测？

Posted 2023-03-27

【中文标题】如何从 pytorch 模型并行化模型预测？

【英文标题】：How to parallelize model prediction from a pytorch model?

【发布时间】：2020-03-08 11:19:50

【问题描述】：

我有一个 PyTorch 实例分割模型，可以一一预测图像的掩码。有什么办法可以并行化这个任务，以便它可以更快地完成。我曾尝试使用多处理pool.apply_async() 来调用执行预测并传递所需参数的方法，但它会引发分段错误。任何帮助深表感谢。谢谢。

【参考方案1】：

一般来说，您应该能够使用torch.stack 将多个图像堆叠在一起，然后将其提供给您的模型。不过，如果没有看到你的模型，我不能肯定地说。 （即，如果您的模型是为一次显式处理一张图像而构建的，则这将不起作用）

model = ... # Load your model

input1 = torch.randn(3, 32, 32)  # An RGB input image
input2 = torch.randn(3, 32, 32)  # An RGB input image 

model_input = torch.stack((input1, input2))
print(model_input.shape)         # (2, 3, 32, 32)

result = model(model_input)
print(result.shape)              # (2, 1, 32, 32) 2 output masks, one for each image

如果您自己训练过模型，这看起来很熟悉，因为它本质上是我们在训练期间向网络提供批量图像的方式。

您可以将两个以上堆叠在一起，这通常会受到您可用的 GPU 内存量的限制。