Follow-up question #2150 about how to define an "epoch" when the data loader is infinite

[gabolsgabs]

Hello,

Following up the issue #2150 about how to define an "epoch" when the data loader is infinite, setting __len__ indeed does the trick but it also restarts the infinite IterableDataset at the end of each epoch.

For illustrating the problem, let's imagine a dataset from which we will get information from a pseudo-infinite number of files:

class InfiniteDataset(torch.utils.data.IterableDataset):
    def __init__(self, n_files, n_steps):
       self.n_files = n_files
       self.n_steps = n_steps    # defines an epoch -> thus how often we check the callbacks, compute the val_loss, etc.

    def __iter__(self):
       while True:
           for x in range(self.n_files):
               yield x    # this correspond to the index file we would like to work with and do whatever

    def __len__(self):
       return self.n_steps

In the first epoch, I expect to use the files from [:n_steps] and then, for the next epochs the files from [n_steps*current_epoch:n_steps*current_epoch+n_step]. However at each epoch I only have access to the [:n_steps] files. Here's a running example:

class DumpModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.automatic_optimization = False
        self.some_weights = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))


    def training_step(self, batch, batch_idx):
        print(batch)    # check the info given by the dataloader
        return

    def validation_step(self, batch, batch_idx):
        return

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

ds_train = InfiniteDataset(n_files=999, n_steps=3)
model = DumpModel()
trainer = pl.Trainer(max_epochs=3)
trainer.fit(model, train_dataloaders=ds_train)

Outputs

Epoch 0:   0%|                                                                                                                                               | 0/3 [00:00<00:00, 4604.07it/s]
0    # expected file index  
Epoch 0:  33%|████████████████████████████████████████████▋                                                                                         | 1/3 [00:00<00:00, 2099.25it/s, v_num=0]
1
Epoch 0:  67%|█████████████████████████████████████████████████████████████████████████████████████████▎                                            | 2/3 [00:00<00:00, 1800.39it/s, v_num=0]
2
Epoch 1:   0%|                                                                                                                                      | 0/3 [00:00<00:00, 1973.79it/s, v_num=0]
0   # expected file index -> should be 3
Epoch 1:  33%|████████████████████████████████████████████▋                                                                                         | 1/3 [00:00<00:00, 1093.41it/s, v_num=0]
1   # expected file index -> should be 4
Epoch 1:  67%|██████████████████████████████████████████████████████████████████████████████████████████                                             | 2/3 [00:00<00:00, 985.20it/s, v_num=0]
2   # expected file index -> should be 5
Epoch 2:   0%|                                                                                                                                      | 0/3 [00:00<00:00, 2790.62it/s, v_num=0]
0   # expected file index -> should be 6
Epoch 2:  33%|████████████████████████████████████████████▋                                                                                         | 1/3 [00:00<00:00, 1700.16it/s, v_num=0]
1   # expected file index -> should be 7
Epoch 2:  67%|█████████████████████████████████████████████████████████████████████████████████████████▎                                            | 2/3 [00:00<00:00, 1552.68it/s, v_num=0]
2   # expected file index -> should be 8
Epoch 2: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 53.14it/s, v_num=0]

Is there an easy way of defining the n_steps of each epoch for an IterableDataset without restarting it? i.e. to be able to replicate this behavior:

ds_train_iter = ds_train.__iter__()

for epoch in range(0, n_epochs):
    for idx in range(0, n_steps):
         batch=next(ds_train_iter)

Thanks!