PyTorch - 如何並行處理 model 輸入？

Question

考慮以下簡單的神經網絡：

class CustomNN(torch.nn.Module):
    def __init__(self):
        super(CustomNN, self).__init__()

    def forward(self, x):
        sleep(1)
        return x

我想知道我們是否可以並行調用forward() 。 按照官方教程，我認為以下代碼可以工作：

x = torch.rand(10, 5).cuda()
futures = [torch.jit.fork(model, x[i,:]) for i in range(10)]
results = [torch.jit.wait(fut) for fut in futures]

我預計它會在大約 1 秒內運行，但它仍然會休眠整整 10 秒。 有什么方法可以並行調用 model 嗎？

Answer 1

不幸的是，您不能在 python 中執行此操作，根據torch.jit.fork要為 TorchScript 模塊提供服務，您需要使用適當的線程池制作 C++ 應用程序 -請參見此處

此外，您的代碼不包括（或您不顯示）轉換。 您必須使用torch.jit.script包裝模塊，以便將其編寫（或跟蹤）到 ScriptModule

traced_NN = torch.jit.script(CustomNN())

即使如此，它也不起作用，因為 TorchScript 僅支持 pytorch 函數（甚至不完全）、python 內置函數和math模塊（請參閱此處）