Pre-computing total number of training steps
🚀 Feature
As it stands, Lightning does not expose the anticipated total number of training steps for a given run. For instance, let's say we specify max_epochs as 100. Our data loader has a total length of 24, our batch size is 4, and we use 2 GPUs. From this we could easily calculate that we'll perform 100 * 24 / (2 * 4) = 300 optimizer steps. It would be convenient to expose this directly, e.g. as a module attribute like self.total_steps.
Given that training duration can be specified either via max_steps or max_epochs, some decision logic is required.
Motivation
Some LR schedulers need to know in advance how many steps will be taken by the optimizer. For instance, if we do linear LR decay to zero over the full duration of training, we need to know this number ahead of time to calculate an appropriate per-step decay factor (example here: https://huggingface.co/transformers/_modules/transformers/optimization.html#get_linear_schedule_with_warmup).
Alternatives
We're currently calculating this value manually and so are some other configurations out there (example here: https://github.com/nateraw/hf-text-classification/blob/ef4caef22239c9988e4a704512d53bd629d52e8f/train.py#L38). However, it's tedious boilerplate that would be nice to abstract away.
Additional context
I've discussed this on Slack: https://pytorch-lightning.slack.com/archives/CRBLFHY79/p1603786080478900 The upshot was that it's currently not available but could be of interest. To track this properly, I'm adding a feature request here.