How does one define an "epoch" when the data loader is infinite?

[hbredin]

I just started playing with pytorch-lightning API to decide whether to make the switch for my own speech processing project.

I have an infinite DataLoader that wraps an IterableDataset that basically extracts random audio chunks from a large speech corpus, forever. As such, it has no __len__ and never raises StopIteration exception.

Therefore, calling trainer.fit(...) stays at Epoch 1 forever as the batch number keeps increasing.

I am looking for a way to tell trainer.fit(...) that it should consider that an epoch is completed every N batches. Is there any?

[github-actions]

Hi! thanks for your contribution!, great first issue!

[s-rog]

set __len__ to return n

Edit:
when working with datasets of an arbitrary length I usually have n set by hparams instead of being hardcoded

[hbredin]

Thanks for your quick answer, @s-rog. It does the trick.

[gabolsgabs]

Hello,

I have a follow-up question regarding this solution. 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?

Thanks!

[AlexHarn]

I am also currently looking into this as I am facing a similar issue in my own project.

As far as I understand, the data loader is not actually reset/restarted between each epoch. There is an option for that (reload_dataloaders_every_n_epochs (int) in the trainer), but it is not set per default.

However, the iteration is restarted for each epoch, so a new iterator is created by calling __iter__ of the IterableDataset again, and this is where you are misunderstanding what's happening. You need to keep track of where you stopped during the last training iteration yourself inside of your InfiniteDataset implementation. An intuitive way in your minimal example might look like this:

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.
       self.current_file = 0

    def __iter__(self):
       while True:
           yield self.current_file
           self.current_file += 1
           if self.current_file >= self.n_files:
               self.current_file = 0

    def __len__(self):
       return self.n_steps

I personally prefer to manually handle the iteration instead of using yield though, especially because it allows me to properly implement an iterable dataset that can leverage multiple data loader worker threads. In your minimal example, that would look like this:

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.
       self.current_file = 0

    def __iter__(self):
        self.current_file -= 1
        return self

    def __next__(self):
        self.current_file += 1
        if self.current_file >= self.n_files:
            self.current_file = 0
        return self.current_file

    def __len__(self):
       return self.n_steps

[thongonary]

I came across this thread as I'm working on something similar.

While @AlexHarn's solution is correct, I would like to add that if you want to use multiple workers in DataLoader (num_workers > 1), you have to set persistent_workers=True so that the worker processes will not be killed after each epoch and the object states are maintained.

[AlexHarn]

I came across this thread as I'm working on something similar.

While @AlexHarn's solution is correct, I would like to add that if you want to use multiple workers in DataLoader (num_workers > 1), you have to set persistent_workers=True so that the worker processes will not be killed after each epoch and the object states are maintained.

Thanks for adding that, it's very important when doing something like this and forgetting about it can cause a lot of issues.