How to set scheduler warmup steps according to warmup ratio with respect to total training step properly?

[BuxianChen]

I am organizing my pure pytorch code to pytorch-lightning, but I can't find an elegant way to set scheduler like this, is there anyone do me a favor?

import pytorch_lightning as pl
import torch
from transformers import get_linear_schedule_with_warmup

class ModelPL(pl.LightningModule):
    def __init__(self, model, config):
        super().__init__()
        self.model = model
        self.config = config

    # training_step, train_dataloader, etc. omit here

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(),
            lr=0.001,
            )
        # ===== how to get these parameters properly ? ======
        # num_training_steps = len(train_dataset) * total_epochs / accumulate_grad_batches
        # num_warmup_steps = num_training_steps * warmup_ratio
 
        scheduler = get_linear_schedule_with_warmup(
            optimizer=optimizer,
            num_warmup_steps=500,  # TO REPLACE
            num_training_steps=10000  # TO REPLACE
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1,
            }
        }

model_config = {
    "scheduler": {
        "warmup_ratio": 0.05
    }
}

train_config = {
    "max_epochs": 2,
    "accumulate_grad_batches": 4
}

# model: pytorch Module
model_pl = ModelPL(model, model_config)
trainer = pl.Trainer(**train_config)
trainer.fit(model_pl)

[ir2718]

In case anyone runs into this problem, my solution was the following:

1. In your pl.LightningDataModule constructor, define a variable that calculates the length of the dataloader
2. Create a singleton class DataLoaderinfo and set the value of the dataloader size to the calculated value
3. Use the getter of the singleton class to get the length of the dataloader and calculate the number of warmup steps and training steps dynamically

The solution is kinda funky because you need to initialize the dataset object without actually using it, but I couldn't think of a better way. I would also like to mention that I'm using the LightningCLI class, and things work as expected without the need of overriding configure_optimizers.