Link arguments from Datamodule into init_args of lr_scheduler
Hey!
I'm trying to use LightningArgumentParser.link_arguments to link an argument from the Datamodule to the init_args of the LR scheduler with:
class MyLightningCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        # Set lr_scheduler's num_training_steps from datamodule class
        parser.link_arguments(
            "data",
            "lr_scheduler.init_args.num_training_steps",
            compute_fn=lambda dm: dm.get_num_training_steps(),
            apply_on="instantiate",
        )
and I get the following error:
ValueError: No action for key "lr_scheduler.init_args.num_training_steps".
I was wondering if such thing is possible, or is linking to init_args is exclusive to the model and data classes?

  Code to reproduce
trainer.py
import pytorch_lightning as pl
import torch.nn
from pytorch_lightning.utilities.cli import LR_SCHEDULER_REGISTRY
from pytorch_lightning.utilities.cli import LightningCLI
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader, Dataset


class MyLightningCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        # Set lr_scheduler's num_training_steps from datamodule class
        parser.link_arguments(
            "data",
            "lr_scheduler.init_args.num_training_steps",
            compute_fn=lambda dm: dm.get_num_training_steps(),
            apply_on="instantiate",
        )


@LR_SCHEDULER_REGISTRY
class WarmupLR(LambdaLR):
    def __init__(
        self,
        optimizer: Optimizer,
        warmup_proportion: float,
        num_training_steps: int,
        last_epoch=-1,
    ) -> None:
        self.num_training_steps = num_training_steps
        self.num_warmup_steps = round(num_training_steps * warmup_proportion)
        super().__init__(optimizer, lr_lambda=self.lr_lambda, last_epoch=last_epoch)

    def lr_lambda(self, current_step: int) -> float:
        if current_step < self.num_warmup_steps:
            return float(current_step) / float(max(1, self.num_warmup_steps))
        return max(
            0.0,
            float(self.num_training_steps - current_step)
            / float(max(1, self.num_training_steps - self.num_warmup_steps)),
        )


class DataModule(pl.LightningDataModule):
    def __init__(self, name):
        super().__init__()
        self.length = len(name)

    def train_dataloader(self):
        return DataLoader(Dataset())

    def get_num_training_steps(self) -> int:
        return self.length


class LitModel(pl.LightningModule):
    def __init__(self, num_labels):
        super().__init__()
        self.num_labels = num_labels
        self.nn = torch.nn.Linear(num_labels, num_labels)

    def training_step(self, *args, **kwargs):
        return


if __name__ == "__main__":
    cli = MyLightningCLI(
        model_class=LitModel,
        datamodule_class=DataModule,
    )
config.yaml
data:
  name: blablabla

model:
  num_labels: 5

optimizer:
  class_path: torch.optim.Adam
  init_args:
    lr: 0.01

lr_scheduler:
  warmup_proportion: 0.1

trainer:
  max_epochs: 2
You need to add an empty configure_optimizers method to your model as there's a bug that disallows leaving it unimplemented. It will be fixed with #11672
The error appears because the lr_scheduler arguments have not been added yet. You can see the order here:

  
    
      pytorch-lightning/pytorch_lightning/utilities/cli.py
    
    
        Lines 603 to 609
      in
      86b177e
    
  
  
    

        
          
           self.add_arguments_to_parser(parser) 
        

        
          
           # add default optimizer args if necessary 
        

        
          
           if not parser._optimizers:  # already added by the user in `add_arguments_to_parser` 
        

        
          
               parser.add_optimizer_args(OPTIMIZER_REGISTRY.classes) 
        

        
          
           if not parser._lr_schedulers:  # already added by the user in `add_arguments_to_parser` 
        

        
          
               parser.add_lr_scheduler_args(LR_SCHEDULER_REGISTRY.classes) 
        

        
          
           self.link_optimizers_and_lr_schedulers(parser) 
        
    
  


So you have two options:

Delay the linking until we've automatically added the lr scheduler classes

class MyLightningCLI(LightningCLI):
    @staticmethod
    def link_optimizers_and_lr_schedulers(parser):
        # Set lr_scheduler's num_training_steps from datamodule class
        parser.link_arguments(
            "data",
            "lr_scheduler.init_args.num_training_steps",
            compute_fn=lambda dm: dm.get_num_training_steps(),
            apply_on="instantiate",
        )
        LightningCLI.link_optimizers_and_lr_schedulers(parser)

Manually add the classes yourself at this hook:

class MyLightningCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        # Manually add the lr scheduler classes
        parser.add_lr_scheduler_args(LR_SCHEDULER_REGISTRY.classes)
        # Set lr_scheduler's num_training_steps from datamodule class
        parser.link_arguments(
            "data",
            "lr_scheduler.init_args.num_training_steps",
            compute_fn=lambda dm: dm.get_num_training_steps(),
            apply_on="instantiate",
        )
Also, the config for the scheduler should be:
lr_scheduler:
  class_path: __main__.WarmupLR
  init_args:
    warmup_proportion: 0.01
    num_training_steps: 1