Многопроцессорная функция внутри Pytorch вперед - PullRequest
Новая
       6

Многопроцессорная функция внутри Pytorch вперед

1 голос
/ 12 марта 2020

Я пытаюсь выполнить полностью подключенный слой к каждому узлу, внедряющемуся после конституционного уровня графа. Это работает, но очень медленно, поэтому я пытаюсь его парализовать.

В основном у меня есть этот код для нормальной версии: 

class Net(nn.Module):
def __init__(self):
    super(Net, self).__init__()
    # first 3 is the feature size
    self.gcn = nn.ModuleList([
        GCN(3, 32, F.leaky_relu),
        GCN(32, 32, F.leaky_relu),
        GCN(32, 32, F.leaky_relu),
        GCN(32, 32, F.leaky_relu),
        GCN(32, 32, F.leaky_relu),
        GCN(32, 32, F.leaky_relu),
        GCN(32, 32, F.leaky_relu),
        GCN(32, 32, F.leaky_relu),
        GCN(32, 32, F.leaky_relu),
        GCN(32, 32, None)])

    self.fc = nn.ModuleList([
        nn.Linear(32, 32, bias=True),
        nn.Linear(32, 32, bias=True),
        nn.Linear(32, 1, bias=True)])


def forward(self, g, features):

    ret = th.zeros(len(features), 1, dtype=features.dtype, device=features.device)

    for layer in self.gcn:
        features = layer(g, features)

    for i, f in enumerate(features):
        for l in self.fc:
            f = l(f)
        ret[i][0] = f

    return ret

в теории второе значение должно легко распараллеливаться. Проблема в том, что я не смог этого сделать. Я попробовал torch.multiprocessing с чем-то вроде 

def tail(layers, i, f, dictionary):
    for l in layers:
        f = l(f)
    dictionary[i]: f

def forward(self, g, features):

    ret = th.zeros(len(features), 1, dtype=features.dtype, device=features.device)

    for layer in self.gcn:
        features = layer(g, features)

    manager = mp.Manager()
    return_dict = manager.dict()

    proc = []
    for i, f in enumerate(features):
        p = mp.Process(target=tail, args=(self.fc, i, f, return_dict))
        p.start()
        proc.append(p)

    for p in proc:
        p.join()

    for i in return_dict.keys():
        ret[i] = return_dict[i]

    return ret

Но не повезло. Есть идеи?

...