Есть ли в Pytorch-Lightning модуль многопроцессорности (или Joblib)?

Question

Я искал в Google, но не могу найти, есть ли multiprocessing модуль, доступный в Pytorch-Lightning, точно так же, как в Pytorch есть torch.multiprocessing модуль.

Кто-нибудь знает, есть ли у Pytorch-Lightning это (или Joblibаналогичный) модуль? Я ищу модуль Pytorch-Lightning, который позволяет мне распараллеливать несколько графических процессоров.

Спасибо заранее.

Изменить: Чтобы быть более конкретным, я ищуmultiprocessing модуль в Pytorch-Lightning, который позволяет мне распараллеливать несколько графических процессоров для вычислений, не связанных с нейронной сетью, таких как:

import numpy as np
import torch
from torch.multiprocessing import Pool

X = np.array([[1, 3, 2, 3], [2, 3, 5, 6], [1, 2, 3, 4]])
X = torch.DoubleTensor(X)

def X_power_func(j):
    X_power = X.cuda()**j
    return X_power

if __name__ == '__main__':
  with Pool(processes = 2) as p:   # Parallelizing over 2 GPUs
    results = p.map(X_power_func, range(4))

results

Answer 1

Да, в основном все, что вам нужно сделать, это предоставить Trainer с соответствующими аргументами gpus=N и укажите backend:

# train on 8 GPUs (same machine (ie: node))
trainer = Trainer(gpus=8, distributed_backend='ddp')

# train on 32 GPUs (4 nodes)
trainer = Trainer(gpus=8, distributed_backend='ddp', num_nodes=4)

Вы можете узнать больше об этом в учебной документации для нескольких GPU.

РЕДАКТИРОВАТЬ:

На самом деле вы искали distributed модуль вместо multiprocessing, torch.distributed.DistributedDataParallel обычно рекомендуется для распараллеливания на нескольких графических процессорах.