Support A Variable Number of Batches

[e-yi]

Description & Motivation

I have a customized batch sampler that has an undetermined number of batches (to maximize the use of GPU memory with variable-sized samples). I believe this behavior is supported by vanilla PyTorch. However, in Lightning, the number of batches is precalculated, as shown here:

pytorch-lightning/src/lightning/pytorch/loops/fit_loop.py

Line 254 in 8ad3e29

num_batches = _parse_num_batches(stage, length, trainer.limit_train_batches)

Pitch

Support a variable number of batches by calculating the batch limit at the beginning of each epoch or at the end of last epoch.

Alternatives

No response

Additional context

No response

cc @lantiga @Borda @tchaton

[e-yi]

For anyone facing the same issue, a simple temporary workaround is to return float('inf') in your BatchSampler.__len__ method.

[dadwadw233]

I meet the same problem while i trying to traning my model with dynamic batch sequence length setting and i have a same strategy with @e-yi to maximize the use of GPU, my program looks like running correctly until i find it skip the epoch end validation 😭 . i check the lightning.loops.fetchers and find:

@override
    def __next__(self) -> _ITERATOR_RETURN:
        assert self.iterator is not None
        self._start_profiler()
        try:
            batch = next(self.iterator)
        except StopIteration:
            self.done = True
            raise
        finally:
            self._stop_profiler()
        self.fetched += 1
        if self.length is not None:
            self.done = self.fetched >= self.length
        return batch

‘self.fetched += 1’ lead to this issue.
Although maybe this not handled situation can be ignored by set val_check_interval==step or as @e-yi suggestion, I'm curious if there is a more elegant solution.
@lantiga