How to iterate over Dataloader until a number of samples is seen?

user_10787371 0
Viewed 2452

I'm learning pytorch, and I'm trying to implement a paper about the progressive growing of GANs. The authors train the networks on the given number of images, instead of for a given number of epochs.

My question is: is there a way to do this in pytorch, using default DataLoaders? I'd like to do something like:

loader = Dataloader(..., total=800000)
for batch in iter(loader):
   ... #do training

And the loader loops itself automatically until 800000 samples are seen.

I think that I'd be a better way, than to calculate the number of times you have to loop through the dataset by yourself

dataloader pytorch deep-learning python
1 Answers

You can use torch.utils.data.RandomSampler and sample from your dataset. Here is a minimal setup example:

class DS(Dataset):
    def __len__(self):
        return 5
    def __getitem__(self, index):
        return torch.empty(1).fill_(index)

>>> ds = DS()

Initialize a random sampler providing num_samples and setting replacement to True i.e. the sampler is forced to draw instances multiple times if len(ds) < num_samples:

>>> sampler = RandomSampler(ds, replacement=True, num_samples=10)

Then plug this sampler to a new torch.utils.data.DataLoader:

>>> dl = DataLoader(ds, sampler=sampler, batch_size=2)

>>> for batch in dl:
...     print(batch)
tensor([[6.],
        [4.]])
tensor([[9.],
        [2.]])
tensor([[9.],
        [2.]])
tensor([[6.],
        [2.]])
tensor([[0.],
        [9.]])
user_63313690
Related