How to share Dataset among different processes in Pytorch？

distributed
1

My dataset for training Pytorch model needs plenty of I/O to generate, so I hope the program will use only 1 process to generate it, and this is my code:

import torch.multiprocessing as mp

# world_size is GPUs' num
mp.spawn(train_worker, args=(world_size, config_data), nprocs=world_size, join=True)

def train_worker(rank, world_size, cfgs):
    trainer = MyDistributedTask(rank, world_size, cfgs)
    trainer.run()

class MyDistributedTask():
    def make_dataset(self, is_tra_val_te: str):
        # Use process 0 to generate dataset
        if self.rank == 0:
            data_cfgs = self.cfgs["data_cfgs"]
            dataset = MyDataset(data_cfgs, is_tra_val_te)
            return dataset

However, MyDataset generated in process 0 is not sent to other processes, so how can I share MyDataset among processes? And I don’t want to use mp.Lock because it will make program slower and make polars appear deadlock mistakes.

Thank for your reply.