Does Pytorch-Lightning have a multiprocessing (or Joblib) module?

[leockl]

❓ Questions and Help

What is your question?

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.

Ps. Sorry if this this the wrong place to post this question. I have posted the same question in Stackoverflow, but haven't received a reply.

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

[Borda]

PL is just PyTorch under the hood, you can use torch on joblib directly...
in case you want train distributed in CPU only you can use ddp_cpu backend 🐰