Configuring OneCycleLR from yaml file lightning CLI

[zarkoivkovicc]

Is there a way to configure OneCycleLR using yaml config files and Lightning CLI?
The problem is that the argument of OneCycleLR on initialization is the total number of steps, which I usually initialize using self.trainer.estimated.stepping.batches inside configure_optimizers inside lightning module. I don't see how this could be done using CLI and config files.

For the reference, I implemented CLI as described here

[mauvilsa]

It would be nice if this were possible by doing:

class MyCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        parser.link_arguments(
            "trainer.estimated_stepping_batches",
            "model.scheduler.init_args.total_steps",
            apply_on="instantiate",
        )

Unfortunately this is not possible because the trainer is instantiated by the LightningCLI class and not by jsonargparse, which prevents linking with source the trainer. There are other use cases which could be solved by linking from the trainer but not possible because of this. I have thought about refactoring LightningCLI to allow it. But this isn't a simple task and haven't had time for it.

Though, I think there could be a non-optimal workaround. Something like:

class MyModel(LightningModule):
    def __init__(
        self,
        optimizer: OptimizerCallable = torch.optim.Adam,
        scheduler: LRSchedulerCallable = torch.optim.lr_scheduler.ConstantLR,
    ):
        super().__init__()
        self.optimizer = optimizer
        self.scheduler = scheduler

    def configure_optimizers(self):
        optimizer = self.optimizer(self.parameters())
        scheduler = self.scheduler(optimizer)
        if isinstance(scheduler, torch.optim.lr_scheduler.OneCycleLR):
            scheduler.total_steps = self.trainer.estimated_stepping_batches
        return {"optimizer": optimizer, "lr_scheduler": scheduler}

if __name__ == "__main__":
    cli = LightningCLI(MyModel, auto_configure_optimizers=False)

Note that I haven't tested it. It is just to illustrate the idea.

Not optimal because when wanting to use OneCycleLR, it would be required to specify a dummy total_steps in the config file, which would then be overwritten in configure_optimizers.

[glsch]

I also noticed that if one wants multiple optimizers and/or schedulers and writes, say,

rnn_optimizer: OptimizerCallable = torch.optim.Adagrad

instead of

rnn_optimizer: OptimizerCallable = lambda p: torch.optim.Adagrad(p)

init_args of the rnn_optimizer are not injected in the configuration.

Even with rnn_optimizer: OptimizerCallable = torch.optim.Adagrad, I would expect this:

  rnn_optimizer:
    class_path: torch.optim.Adagrad
    init_args:
      lr: 0.01
      lr_decay: 0.0
      weight_decay: 0.0
      initial_accumulator_value: 0.0
      eps: 1.0e-10
      foreach: null
      maximize: false
      differentiable: false

but get this:

  rnn_optimizer: torch.optim.Adagrad

It is also not possible to set defaults in the CLI for when I am using lambda.

The following code:

class RnnClusterer(LightningModule):
    def __init__(self,
                 rnn: torch.nn.Module,
                 noise: torch.nn.Module,
                 head: torch.nn.Module,
                 main_heads: int,
                 control_heads: int,
                 rnn_optimizer: OptimizerCallable = lambda p: torch.optim.Adagrad(p),
                 head_optimizer: OptimizerCallable = lambda p: torch.optim.Adagrad(p),
                 rnn_lr_scheduler: LRSchedulerCallable = lambda o: torch.optim.lr_scheduler.LinearLR(o),
                 head_lr_scheduler: LRSchedulerCallable = lambda o: torch.optim.lr_scheduler.LinearLR(o),
                 scheduler_config: Optional[dict] = None,
                 leak: Union[float, None, str] = "auto",
                 accumulation_steps: int = 1,
                 leak_decay: float = 0.5,
                 in_model_logging_level: int = logging.INFO,
                 clusteriness: float = 0.8,
                 ):
        super().__init__()

with this custom CLI:

parser.set_defaults({
            "model.rnn_lr_scheduler.init_args.start_factor": 1,
            "model.rnn_lr_scheduler.init_args.end_factor": 0.0,
            "model.rnn_lr_scheduler.init_args.total_iters": 1000,
            "model.head_lr_scheduler.init_args.start_factor": 1,
            "model.head_lr_scheduler.init_args.end_factor": 0.0,
            "model.head_lr_scheduler.init_args.total_iters": 1000
        })

Produces the following error:

jsonargparse._namespace.NSKeyError: No action for key "model.rnn_lr_scheduler.init_args.start_factor" to set its default.