TODO tracker for ongoing training loop refactors #7938
Closed
Description
📚 TODOs around the loop refactor
This issue tracks the work planned for and after the loop refactor #7700.
Core contributors feel free to update the description here directly.
Hygiene and code quality
- Typing for full mypy compatibility (see ignored files in setup.cfg)
- Documentation for the public interface + address [Refactor] Remove _run_evaluation + 3 EvaluationLoop #8065 (comment)
- Evaluate which methods should be public and which should be protected
- Rename Trainer.train_loop to Trainer.fit_loop and deprecate the former access.
- Formalize how arguments are forwarded to lower loops (especially in
connect
Feature integration and follow-up refactors
- Extract a OptimizerLoop, potentially a loop for TBPTT extract optimizer loop #9191
- extract utility for
_run_optimization
- for manual optimization, replace closure call with direct call to
_training_step()
- extract utility for
- Integrate progress tracking Formalize progress tracking inside of the trainer internals #6429, this will also increase clarity and minimize attack surface for bugs!
- Determine who should own the dataloaders and responsibility of resetting them Loop Refactor 1/N - Training Loop #7871 (comment). Can references fall out of sync? Loop Refactor 1/N - Training Loop #7871 (comment)
- Add a teardown, let loops move their results to CPU instead of LoggerConnector: Move result teardown to loops #8245
- The
EvaluationDataLoaderLoop.predictions
are currently accumulated inside the inner evaluation loop instead of per epoch and then combined over all dataloaders. Remove write_predictions from LightningModule #8850 - Rename inner loops to epoch_loop, batch_loope etc to distinguish top level loop access. Also propose file structure organizing these. Resolved by
pytorch_lightning.loops
file structure: group by dataloader, epoch, and batch loop #8077
Customization
- Formalize how loops get instantiated, set on the trainer, and connected to each other. Loop specialization #8226
Miscellaneous
- Several todo's are annotated as TODO(@awaelchli)
- Some todo's are annotated as TODO(@justusschock) as well.