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

Question

我有一个PyTorch实例分割模型，它一个接一个地预测图像的掩模。有没有办法将这个任务并行化，这样它就可以更快地完成。我曾尝试使用多处理pool.apply_async()来调用传递所需参数的预测方法，但它抛出了分段错误。任何帮助都是非常感谢的。谢谢。

Answer 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内存量的限制。