Lightning-AI / Lightning-AI/pytorch-lightning
Support A Variable Number of Batches
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### 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: https://github.com/Lightning-AI/pytorch-lightning/blob/8ad3e29816a63d8ce5c00ac104b14729a4176f4f/src/lightning/pytorch/loops/fit_loop.py#L254
### 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
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading the batch-limit calculation in src/lightning/pytorch/loops/fit_loop.py at the linked location, then trace how the customized batch sampler is expected to report its batches. Define completion as calculating the batch limit at the beginning of each epoch or from the previous epoch while preserving training-loop behavior; the issue provides no test path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100