Как одновременно запустить несколько веток в pytorch?

Question

Я пытался построить сеть с несколькими ветвями в pytorch . Но как я могу запускать несколько ветвей параллельно вместо , запускать их одну за другой ?

В отличие от tenorflow или keras, pytorch использует динамический график c, поэтому я не могу заранее определить параллельную обработку.

Я искал какую-то аналогичную официальную реализацию сети pytorch, такую ​​как Inceptio nNet, только чтобы узнать, 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)

, вероятно, выполняются более или менее одновременно. То же самое для всех следующих слоев.