PyTorch alternative for tf.data.experimental.sample_from_datasets

user_987397 0
Viewed 828

Suppose I have two datasets, dataset one with 100 items and dataset two with 5000 items.

Now I want that during training my model sees as much items from dataset one as from dataset two.

In Tensorflow I can do:

dataset = tf.data.experimental.sample_from_datasets(
    [dataset_one, dataset_two], weights=[50,1], seed=None
)

Is there an alternative in PyTorch that does the same?

I think this is not too difficult to implement by creating a custom dataset (not working example)

from torch.utils.data import Dataset

class SampleDataset(Dataset):
    def __init__(self, datasets, weights):
        self.datasets = datasets
        self.weights = weights

    def __len__(self):
        return sum([len(dataset) for dataset in self.datasets])

    def __getitem__(self, idx):
        # sample a random number and based on that sample an item
        
        return self.datasets[dataset_idx][sample_idx]

However, this seems quite common. Is there already something like this available?

pytorch tensorflow
1 Answers

I don't think there is a direct equivalent in PyTorch.

However, there's a function called torch.utils.data.WeightedRandomSampler which samples indices based on a list of probabilities. You can use this in combination with torch.data.utils.ConcatDataset and torch.utils.data.DataLoader's sampler option.

I'll give an example with two datasets: SetA has 500 elements and SetB which only has 10.

First, you can create a concatenation of all your datasets with ConcaDataset:

ds = ConcatDataset([SetA(), SetB()])

Then, we need to sample it. The problem is, you can't just give WeightedRandomSampler [50, 1], as you did in Tensorflow. As a workaround, you can create a list of probabilities of the same length as the size of the total dataset.

The corresponding probability list for this example would be:

dist = np.array([1/51]*500 + [50/51]*10)

Essentially, the first 500 indices (i.e. indices 'pointing' to SetA) will have a probability of 1/51 of being choosen while the following 10 indices (i.e. indices in SetB) will have a probability of 50/51 (i.e much more likely to being sampled since there are less elements in SetB, this is the desired result!)

We can create a sampler from that distribution:

WeightedRandomSampler(dist, 10)

Where 10 is the number of sampled elements. I would put the size of the smallest dataset, otherwise you would likely be going over the same datapoints multiple times during the same epoch...

Finally, we just have to instanciate the dataloader with our dataset and sampler:

dl = DataLoader(ds, sampler=sampler)

To summarize:

ds = ConcatDataset([SetA(), SetB()])
dist = np.array([1/51]*500 + [50/51]*10)
sampler = WeightedRandomSampler(dist, 10)
dl = DataLoader(ds, sampler=sampler)

Edit, for any number of datasets:

sets = [SetA(), SetB(), SetC()]
ds = ConcatDataset(sets)

dist = np.concatenate([[(len(ds) - len(s))/len(ds)]*len(s) for s in sets])
sampler = WeightedRandomSampler(weights=dist, num_samplesmin([len(s) for s in sets])
dl = DataLoader(ds, sampler=sampler)
user_63313690
Related