У меня есть код факела, в котором я пытаюсь вычислить вывод 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
. Я не могу понять, почему это происходит. Может кто-нибудь, пожалуйста, помогите? Заранее спасибо!