Properly support iteration-based training

[Rizhiy]

🚀 Feature

Currently, a lot of training functionality is tied to the concept of "Epoch".
IMHO, training based on epoch is bad practice, and I always train using the raw number of iterations.
PL makes this approach very difficult.

Motivation

Currently, it is very difficult to use PL with iteration based training, e.g. Trainer currently resumes from checkpoint at the start of next epoch, rather than at the iteration on which checkpoint was made.

In the real world, data composition of the training set is changing pretty frequently and training based on epochs can give false results. e.g. if I create a new dataset by duplicating every item in my original, and train for the same number of epochs, I will probably get better result which might make me think that new data I added improved the results, but in reality I just trained for longer.

Training on very large datasets can take days/weeks. When resources are constrained I need to schedule time for each experiment, and having to manually calculate number of epochs that fit in a time period is very annoying.

Pitch

Add proper support for iteration based training, the easiest would be to just keep track of total number of iterations through the epochs and allow various parts/callbacks to use it.

Alternatives

You can also completely change to iteration based training, and calculate epochs using length of the dataset, but that would probably require quite a large rewrite and be about the same.

Additional context

Current problems:

• Not possible to run validation at frequency greater than the length of epoch
• If the dataset is very large and training is only a few epochs, then resume can significantly worsen the results by starting from next epoch
• When max_steps is passed to pl.Trainer would be nice if it was used in tqdm.

[carmocca]

Hi!

Thanks for the thoughtful feature request. This is currently in the roadmap so keep an eye on future releases 😎

[Rizhiy]

I found a workaround for having validation frequency greater than length of epoch and checkpoint resume:
You can create an infinite sampler, which will also make your epoch infinite length, here is an example:

class RepeatSampler:
    def __init__(self, sampler: Sampler):
        self._sampler = sampler

    def __iter__(self):
        while True:
            yield from iter(self._sampler)

    def __len__(self):
        return math.inf

    def __getattr__(self, name: str):
        return getattr(self._sampler, name)

Using this sampler and max_steps in Trainer seems to work, since total number of steps on resume will be correct, only problem is that batch_idx is wrong after resume.

[awaelchli]

@Rizhiy Iteration based training is possible in Lightning by setting max_steps and just ignoring max_epochs.
Properly restoring all state mid-epoch is challenging and missing in Lighting because the dataloader needs to be restored to the state it was in when training was aborted. That is not so trivial but is on the roadmap as @carmocca mentioned.
I would love to have that feature too!

[Rizhiy]

Ok, thanks. I think under most circumstances you only need to restore the sampler, because datasets shouldn't change during training normally.
Assuming sampler is quite light-weight it can just be fast-forwarded, although it will probably require having its own separate random number generator, which is probably not valid in most cases.