python - 如何在pytorch中并发运行多个分支？

Question

我试图在 pytorch 中构建一个具有多个分支的网络。但是我怎样才能并行运行多个分支而不是一个一个地运行它们？

不像tensorflow或keras，pytorch使用动态图，所以我不能事先定义并发处理。

我查找了一些类似于 InceptionNet 的 pytorch 网络官方实现，结果发现 pytorch 连续运行多个分支。

来自 inception.py

def _forward(self, x):
    branch1x1 = self.branch1x1(x)

    branch5x5 = self.branch5x5_1(x)
    branch5x5 = self.branch5x5_2(branch5x5)

    branch3x3dbl = self.branch3x3dbl_1(x)
    branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
    branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)

    branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
    branch_pool = self.branch_pool(branch_pool)

    outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
    return outputs

四个分支一个接一个运行，先是branch1x1，然后是branch5x5，还有branch3x3dbl，branch_pool。然后输出存储它们的结果，稍后将它们连接起来。

岂不是浪费性能？我们该如何处理？

Answer 1

一般情况下，只要使用pytorch提供的函数，就不必关心网络执行的性能。

正如评论中所指出的，所有对 gpu 的调用都是异步的。只要调用不依赖于数据，它就会被执行。所以在你的情况下你有多个分支。 Pytorch 会根据数据的依赖关系来调度所有的操作并执行。由于您不在分支之间共享数据，因此它们将并行执行。

所以在你的情况下

branch3x3dbl = self.branch3x3dbl_1(x)
branch1x1 = self.branch1x1(x)
branch3x3dbl = self.branch3x3dbl_1(x)

可能或多或少同时执行。以下所有层都一样。