Retrieve exact number of training steps in `configure_optimizers`
❓ Questions and Help
What is your question?
Some schedulers may require the total number of training steps. How can I retrieve it?
Code
def configure_optimizers(self) -> (torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR):
    steps = self.trainer.total_training_steps
What have you tried?
I tried computing them from the dataset length and the training parameters. But I don't think this is correct in a multi-gpu, multi-node setting.
def get_total_number_of_steps_for_schedulers(hparams, trainer):
    if trainer.max_steps is not None:
        return trainer.max_steps
    else:
        training_batches_per_epoch = math.ceil(trainer.num_training_batches / (hparams.batch_size * trainer.world_size))
        steps = (hparams.max_epochs * training_batches_per_epoch) // hparams.accumulate_grad_batches
        return steps
What's your environment?
OS Independent
pytorch_lightning: 0.9.0