LightningCLI does not overwrite optimizer defined in LightningModule

[Toekan]

🐛 Bug

Hi,

First of all, thanks for the great work creating pytorch lightning. Lately I've been using the LightningCLI functionality and it really cut down a lot of boilerplate code and made the CLI easier to use. So thanks for that!

Unfortunately I found one detail where I either don't understand the documentation well, or the behaviour is not as intended: the optimizer defined in configuration file does not seem to overwrite one defined in LightningModule. configure_optimizers.

To Reproduce

Initially I used a hardcoded optimizer and scheduler defined in my LightningModule, like:

class MyBoringModel(LightningModule):
  ...
 def configure_optimizers(self):
        optimizer = torch.optim.RMSprop(
            params=[p for p in self.model.parameters() if p.requires_grad],
            lr=self.hparams.initial_lr,
        )

        lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
            optimizer, gamma=self.hparams.scheduler_gamma
        )

        return [optimizer], [lr_scheduler]

Then I decided I wanted to change to a different optimizer, so I added the following to the config I pass in to the LightningCLI:

# my_config.yaml
...

optimizer:
  class_path: torch.optim.SGD
  init_args:
    lr: 0.0001
    momentum: 0.9

...

A dummy example of the script I run:

# cli.py

trainer_defaults = {
        "callbacks": [
            LearningRateMonitor(logging_interval="epoch"),
            # Default latest model checkpoint
            ModelCheckpoint(),
            # Best performing model checkpoint
            ModelCheckpoint(
                filename="best-epoch={epoch}-loss={total_loss/test_epoch:.3f}",
                monitor="total_loss/test_epoch",
                auto_insert_metric_name=False,
            ),
        ]
    }

cli = LightningCLI(
        model_class=MyBoringModel,
        datamodule_class=MyIrrelevantDataModule,
        trainer_defaults=trainer_defaults,
        run=True,
)

When running the code python cli.py fit --config my_config.yaml I get the following warning:

UserWarning: `MyBoringModel.configure_optimizers` will be overridden by `LightningCLI.configure_optimizers`. 

so I, incorrectly, assumed it would take the optimizer from the config, and either take the scheduler from my model's method, or no/a default scheduler. Instead, it seems to have just taken both the optimizer and the scheduler from my model's method. Is this intended? If so, should the documentation of LightningCLI regarding optimizers and schedulers be expanded to clarify this?

Small detail: the reason I assume the optimizer isn't being overwritten, is because LearningRateMonitor is still logging the learning rate of the optimize defined in the LightningModule, so I guess there is a chance the training is happening with the new optimizer, but the callback is logging the wrong optimizer.

Expected behavior

Overwrite the optimizer and scheduler from the LightningModule with whatever is provided from the config file, even if it's only the optimizer, or only the scheduler.

Environment

• PyTorch Lightning Version: 1.6.0
• PyTorch Version: 1.11.0
• Python version: 3.8
• OS (e.g., Linux): Linux
• CUDA/cuDNN version: / (happens both cpu-only runs and gpu runs)
• How you installed PyTorch (conda, pip, source): pip

Thanks!

cc @carmocca @mauvilsa

[carmocca]

I am not able to reproduce the issue:

import torch

from pytorch_lightning.callbacks import LearningRateMonitor
from pytorch_lightning.utilities.cli import LightningCLI
from tests.helpers import BoringModel


class MyBoringModel(BoringModel):
    def configure_optimizers(self):
        optimizer = torch.optim.RMSprop(
            params=[p for p in self.layer.parameters() if p.requires_grad],
            lr=0.1,
        )

        lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.1)

        return [optimizer], [lr_scheduler]


lr_monitor = LearningRateMonitor(logging_interval="epoch")
cli = LightningCLI(model_class=MyBoringModel, run=True, trainer_defaults={"fast_dev_run": 1, "callbacks": lr_monitor})

print(lr_monitor.lrs)
print(cli.trainer.optimizers)
print(cli.trainer.lr_scheduler_configs)

Prints

{'lr-SGD': [0.0001]}                                                                                                                                                                                                                                                                                                                                                                      
[SGD (
Parameter Group 0
    dampening: 0
    foreach: None
    lr: 0.0001
    maximize: False
    momentum: 0.9
    nesterov: False
    weight_decay: 0
)]
[]

which are the values from the config, as expected.