I have been googling around but can't seem to find if there is a multiprocessing
module available in Pytorch-Lightning, just like how Pytorch has a torch.multiprocessing
module.
Does anyone know if Pytorch-Lightning has this (or a Joblib
similar) module? I am looking for a Pytorch-Lightning module which allows me to parallelize over multiple GPUs
Many thanks in advance.
Edit: To be more specific, I am looking for a multiprocessing
module in Pytorch-Lightning which allows me to parallelize over multiple GPUs on non-neural network computations, such as:
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
Yes, basically all you have to do is to provideTrainer
with appropriate argument gpus=N
and specify 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)
You can read more about it in multi-GPU training documentation .
What you were actually looking for is distributed
module instead of multiprocessing
, torch.distributed.DistributedDataParallel
is usually recommended for parallelizing over multiple GPUs.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.