How to divide the dataset when it is distributed

user_12839713 0
Viewed 1312

Now I want to divide a dataset into two parts: the train set and validation set. I know that on a single GPU I can do this using a sampler:

indices = list(range(len(train_data)))
train_loader = torch.utils.data.DataLoader(
      train_data, batch_size=args.batch_size,
      sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:split]),
      pin_memory=True, num_workers=2)

But when I want to train it in a parallel way using torch.distributed, I have to use another sampler, namely, sampler = torch.utils.data.distributed.DistributedSampler(train_data)

So how should I do to use the two samplers, so that I can divide the dataset and distribute it at the same time?

Thank you very much for any help!

pytorch distributed python
1 Answers

You can split torch.utils.data.Dataset before creating torch.utils.data.DataLoader.

Simply use torch.utils.data.random_split like this:

train, validation =
    torch.utils.data.random_split(
        dataset, 
        (len(dataset)-val_length, val_length)
    )

This would give you two separate datasets which could be used with dataloaders however you wish.

user_108864200
Related