Pytorch-Lightning是否具有多处理(或Joblib)模块?

时间:2020-07-27 04:14:28

标签: python multiprocessing pytorch gpu pytorch-lightning

我一直在四处搜寻,但似乎找不到Pytorch-Lightning中是否有multiprocessing模块可用,就像Pytorch具有torch.multiprocessing模块一样。

有人知道Pytorch-Lightning是否具有此(或类似Joblib的模块)吗?我正在寻找一个Pytorch-Lightning模块,该模块可以让我在多个GPU上并行化

非常感谢。

编辑:更具体地说,我正在Pytorch-Lightning中寻找一个multiprocessing模块,该模块可让我在非神经网络计算上并行处理多个GPU,例如:

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

1 个答案:

答案 0 :(得分:1)

是的,基本上您要做的就是为Trainer提供适当的参数gpus=N并指定后端:

# 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)

您可以在multi-GPU training documentation中了解更多信息。

编辑:

您真正要寻找的是distributed模块而不是multiprocessing,通常建议使用torch.distributed.DistributedDataParallel来实现多个GPU的并行化。