How to set lr scheduler from yaml for lightningcli?

[Galaxy-Husky]

I find that the params of optimizer in the yaml file will override configure_optimizers function. But the scheduler will disappear.
Can I override the params of lr scheduler in the configure_optimizers function by yaml file? If so, how to write the code?

from os import path
from typing import Optional, Tuple
import warnings

import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, random_split

from lightning.pytorch import callbacks, LightningDataModule, LightningModule, Trainer
from lightning.pytorch.callbacks import LearningRateMonitor
from lightning.pytorch.cli import LightningCLI
from lightning.pytorch.demos.mnist_datamodule import MNIST
from lightning.pytorch.utilities import rank_zero_only
from lightning.pytorch.utilities.imports import _TORCHVISION_AVAILABLE
from lightning.pytorch.utilities.warnings import PossibleUserWarning
# import logging
#
#
# logging.getLogger("lightning.pytorch").setLevel(logging.INFO)
# logger = logging.getLogger("lightning.pytorch.core")
# logger.addHandler(logging.FileHandler("core.log"))

if _TORCHVISION_AVAILABLE:
    import torchvision
    from torchvision import transforms
    from torchvision.utils import save_image

warnings.filterwarnings("ignore", category=PossibleUserWarning)
DATASETS_PATH = path.join(path.dirname(__file__), "..", "..", "Datasets")


class ImageSampler(callbacks.Callback):
    def __init__(
        self,
        num_samples: int = 3,
        nrow: int = 8,
        padding: int = 2,
        normalize: bool = True,
        norm_range: Optional[Tuple[int, int]] = None,
        scale_each: bool = False,
        pad_value: int = 0,
    ) -> None:

        if not _TORCHVISION_AVAILABLE:  # pragma: no cover
            raise ModuleNotFoundError("You want to use `torchvision` which is not installed yet.")

        super().__init__()
        self.num_samples = num_samples
        self.nrow = nrow
        self.padding = padding
        self.normalize = normalize
        self.norm_range = norm_range
        self.scale_each = scale_each
        self.pad_value = pad_value

    def _to_grid(self, images):
        return torchvision.utils.make_grid(
            tensor=images,
            nrow=self.nrow,
            padding=self.padding,
            normalize=self.normalize,
            value_range=self.norm_range,
            scale_each=self.scale_each,
            pad_value=self.pad_value,
        )

    @rank_zero_only
    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        if not _TORCHVISION_AVAILABLE:
            return

        images, _ = next(iter(DataLoader(trainer.datamodule.mnist_val, batch_size=self.num_samples)))
        images_flattened = images.view(images.size(0), -1)

        # generate images
        with torch.no_grad():
            pl_module.eval()
            images_generated = pl_module(images_flattened.to(pl_module.device))
            pl_module.train()

        if trainer.current_epoch == 0:
            save_image(self._to_grid(images), f"grid_ori_{trainer.current_epoch}.png")
        save_image(self._to_grid(images_generated.reshape(images.shape)), f"grid_generated_{trainer.current_epoch}.png")


class LitAutoEncoder(LightningModule):
    def __init__(self, hidden_dim: int = 64, lr=10e-3):
        super().__init__()
        self.save_hyperparameters()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 3))
        self.decoder = nn.Sequential(nn.Linear(3, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 28 * 28))

    def forward(self, x):
        z = self.encoder(x)
        return self.decoder(z)

    def training_step(self, batch, batch_idx):
        return self._common_step(batch, batch_idx, "train")

    def validation_step(self, batch, batch_idx):
        self._common_step(batch, batch_idx, "val")

    def test_step(self, batch, batch_idx):
        self._common_step(batch, batch_idx, "test")

    def predict_step(self, batch, batch_idx, dataloader_idx=None):
        x = self._prepare_batch(batch)
        return self(x)

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.lr)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler(optimizer, T_max=self.trainer.estimated_stepping_batches),
                "interval": 'step',
                "frequency": 1,
            }
        }

    def _prepare_batch(self, batch):
        x, _ = batch
        return x.view(x.size(0), -1)

    def _common_step(self, batch, batch_idx, stage: str):
        x = self._prepare_batch(batch)
        loss = F.mse_loss(x, self(x))
        if stage == 'train':
            self.log(f"{stage}_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
        else:
            self.log(f"{stage}_loss", loss, on_epoch=True, prog_bar=True, sync_dist=True)
        return loss


class MyDataModule(LightningDataModule):
    def __init__(self, batch_size):
        super().__init__()
        dataset = MNIST(DATASETS_PATH, train=True, download=True, transform=transforms.ToTensor())
        self.mnist_test = MNIST(DATASETS_PATH, train=False, download=True, transform=transforms.ToTensor())
        self.mnist_train, self.mnist_val = random_split(dataset, [55000, 5000])
        self.batch_size = batch_size

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=self.batch_size)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=self.batch_size)

    def predict_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=self.batch_size)


def cli_main():
    lr_monitor = LearningRateMonitor(logging_interval='step')
    trainer_defaults = {"callbacks": [ImageSampler(), lr_monitor],
                        }
    cli = LightningCLI(
        LitAutoEncoder,
        MyDataModule,
        run=False,
        trainer_defaults=trainer_defaults,
        save_config_kwargs={"overwrite": True},
    )
    cli.trainer.fit(cli.model, datamodule=cli.datamodule)


if __name__ == "__main__":
    cli_main()

# lightning.pytorch==2.0.6
seed_everything: 1234
trainer:
  accelerator: gpu
  strategy: auto
  devices: auto
  num_nodes: 1
  precision: 16-mixed
  logger: true
  callbacks: null
  fast_dev_run: false
  max_epochs: null
  min_epochs: null
  max_steps: 1000
  min_steps: null
  max_time: null
  limit_train_batches: null
  limit_val_batches: 2
  limit_test_batches: 2
  limit_predict_batches: 2
  overfit_batches: 0.0
  val_check_interval: null
  check_val_every_n_epoch: 1
  num_sanity_val_steps: 1
  log_every_n_steps: 10
  enable_checkpointing: false
  enable_progress_bar: true
  enable_model_summary: true
  accumulate_grad_batches: 1
  gradient_clip_val: 10
  gradient_clip_algorithm: value
  deterministic: warn
  benchmark: false
  inference_mode: true
  use_distributed_sampler: true
  profiler: null
  detect_anomaly: false
  barebones: false
  plugins: null
  sync_batchnorm: false
  reload_dataloaders_every_n_epochs: 0
  default_root_dir: null
model:
  hidden_dim: 64
  lr: 0.01
data:
  batch_size: 64
optimizer:
  class_path: torch.optim.AdamW
  init_args:
    lr: 1e-4
    betas:
    - 0.9
    - 0.999
    eps: 1.0e-08
    weight_decay: 1e-4
    amsgrad: false
    maximize: false
    foreach: null
    capturable: false
    differentiable: false
    fused: null
lr_scheduler:
  class_path: torch.optim.lr_scheduler.CosineAnnealingLR
  init_args:
    T_max: 100
    eta_min: 0.0
    last_epoch: -1
    verbose: false

[mauvilsa]

The automatic configure_optimizers from LightningCLI only works the simplest case. If you want to use your own configure_optimizers then disable the automatic one, i.e. auto_configure_optimizers=False. If you do this, my recommendation to make optimizers and learning rate schedulers configurable from yaml, is to use dependency injection, as shown in multiple-optimizers-and-schedulers.

[Galaxy-Husky]

Thank you for your advice on using my own optimizer and scheduler.
Btw, I just found that the lr scheduler didn't disappear in the simplest case. It was already configured, and since it stepped by epoch, I didn't see a change in the learning rate. If I change T_max to 2 in the yaml file, I can see a change in the learning rate.