PyTorch

Question

考虑以下简单的神经网络：

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

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

我想知道我们是否可以并行调用 forward()。在 official tutorial 之后，我认为以下代码可以工作：

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 秒。有没有办法并行调用模型？

Answer 1

您不能在 python 中执行此操作，很遗憾，根据 torch.jit.fork 要为 TorchScript 模块提供服务，您需要制作一个具有适当线程池的 C++ 应用程序 – see here

此外，您的代码不包含（或者您不显示）转换。您必须用 torch.jit.script 包装 Module 才能将其编写（或跟踪）到 ScriptModule

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

即便如此，它也不会工作，因为只有 pytorch 函数（甚至不完全），python 内置函数和 math 模块在 TorchScript (see here)

PyTorch - 如何并行处理模型输入？

PyTorch - How to process model input in parallel?

python

parallel-processing