Мультипроцессорная обработка на одном GPU чрезвычайно медленная - PullRequest
0 голосов
/ 16 февраля 2020

У меня есть код факела, в котором я пытаюсь вычислить вывод 10 слоев, которые не зависят друг от друга. Это похоже на случай, когда мультипроцессорная обработка факела идеально подходит. Вот код, который повторяет ситуацию в моем случае:

from torch.multiprocessing import Pool, set_start_method
import torch
from torch.autograd import Variable

from kernel_2_0 import *
from kernel_2_1 import *
from kernel_2_2 import *
from kernel_2_3 import *
from kernel_2_4 import *
from kernel_2_5 import *
from kernel_2_6 import *
from kernel_2_7 import *
from kernel_2_8 import *
from kernel_2_9 import *

def apply_fn(args):
    f = args[0]
    x = args[1]
    return f(x)

def forward_2(x):
    assert isinstance(x.data, torch.cuda.BoolTensor)
    try:
        set_start_method('spawn', force=True)
    except RuntimeError:
        print('Could not set spawn method')
    with Pool(12) as pool:
        mapped_out = pool.map(apply_fn,
                              [(kernel_2_0, x), (kernel_2_1, x), (kernel_2_2, x), (kernel_2_3, x), (kernel_2_4, x),
                               (kernel_2_5, x), (kernel_2_6, x), (kernel_2_7, x), (kernel_2_8, x), (kernel_2_9, x)])

out = Variable(torch.cuda.BoolTensor(x.size()[0], 10))
out[:, 0] = mapped_out[0]
out[:, 1] = mapped_out[1]
out[:, 2] = mapped_out[2]
out[:, 3] = mapped_out[3]
out[:, 4] = mapped_out[4]
out[:, 5] = mapped_out[5]
out[:, 6] = mapped_out[6]
out[:, 7] = mapped_out[7]
out[:, 8] = mapped_out[8]
out[:, 9] = mapped_out[9]

return out

Однако, когда код выполняется, он почти в 100 раз медленнее, чем версия, где я не выполняю мультиобработку, а вместо этого вызываю все ядра. в последовательном порядке. К вашему сведению, я использую torch==1.4.0 и cuda-10.0. Я не могу понять, почему это происходит. Может кто-нибудь, пожалуйста, помогите? Заранее спасибо!

...