Configure hyperparameters from the CLI (Advanced)

Instantiation only mode

The CLI is designed to start fitting with minimal code changes. On class instantiation, the CLI will automatically call the trainer function associated with the subcommand provided, so you don’t have to do it. To avoid this, you can set the following argument:

cli = LightningCLI(MyModel, run=False)  # True by default
# you'll have to call fit yourself:
cli.trainer.fit(cli.model)

In this mode, subcommands are not added to the parser. This can be useful to implement custom logic without having to subclass the CLI, but still, use the CLI’s instantiation and argument parsing capabilities.

Trainer Callbacks and arguments with class type

A very important argument of the Trainer class is the callbacks. In contrast to simpler arguments that take numbers or strings, callbacks expects a list of instances of subclasses of Callback. To specify this kind of argument in a config file, each callback must be given as a dictionary, including a class_path entry with an import path of the class and optionally an init_args entry with arguments to use to instantiate. Therefore, a simple configuration file that defines two callbacks is the following:

trainer:
  callbacks:
    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
      init_args:
        save_weights_only: true
    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
      init_args:
        logging_interval: 'epoch'

Similar to the callbacks, any parameter in Trainer and user extended LightningModule and LightningDataModule classes that have as type hint a class, can be configured the same way using class_path and init_args. If the package that defines a subclass is imported before the LightningCLI class is run, the name can be used instead of the full import path.

From command line the syntax is the following:

$ python ... \
    --trainer.callbacks+={CALLBACK_1_NAME} \
    --trainer.callbacks.{CALLBACK_1_ARGS_1}=... \
    --trainer.callbacks.{CALLBACK_1_ARGS_2}=... \
    ...
    --trainer.callbacks+={CALLBACK_N_NAME} \
    --trainer.callbacks.{CALLBACK_N_ARGS_1}=... \
    ...

Note the use of + to append a new callback to the list and that the init_args are applied to the previous callback appended. Here is an example:

$ python ... \
    --trainer.callbacks+=EarlyStopping \
    --trainer.callbacks.patience=5 \
    --trainer.callbacks+=LearningRateMonitor \
    --trainer.callbacks.logging_interval=epoch

Note

Serialized config files (e.g. --print_config or SaveConfigCallback) always have the full class_path, even when class name shorthand notation is used in the command line or in input config files.

Multiple models and/or datasets

A CLI can be written such that a model and/or a datamodule is specified by an import path and init arguments. For example, with a tool implemented as:

cli = LightningCLI(MyModelBaseClass, MyDataModuleBaseClass, subclass_mode_model=True, subclass_mode_data=True)

A possible config file could be as follows:

model:
  class_path: mycode.mymodels.MyModel
  init_args:
    decoder_layers:
    - 2
    - 4
    encoder_layers: 12
data:
  class_path: mycode.mydatamodules.MyDataModule
  init_args:
    ...
trainer:
  callbacks:
    - class_path: lightning.pytorch.callbacks.EarlyStopping
      init_args:
        patience: 5
    ...

Only model classes that are a subclass of MyModelBaseClass would be allowed, and similarly, only subclasses of MyDataModuleBaseClass. If as base classes LightningModule and LightningDataModule is given, then the CLI would allow any lightning module and data module.

Tip

Note that with the subclass modes, the --help option does not show information for a specific subclass. To get help for a subclass, the options --model.help and --data.help can be used, followed by the desired class path. Similarly, --print_config does not include the settings for a particular subclass. To include them, the class path should be given before the --print_config option. Examples for both help and print config are:

$ python trainer.py fit --model.help mycode.mymodels.MyModel
$ python trainer.py fit --model mycode.mymodels.MyModel --print_config

Models with multiple submodules

Many use cases require to have several modules, each with its own configurable options. One possible way to handle this with LightningCLI is to implement a single module having as init parameters each of the submodules. This is known as dependency injection which is a good approach to improve decoupling in your code base.

Since the init parameters of the model have as a type hint a class, in the configuration, these would be specified with class_path and init_args entries. For instance, a model could be implemented as:

class MyMainModel(LightningModule):
    def __init__(self, encoder: nn.Module, decoder: nn.Module):
        """Example encoder-decoder submodules model

        Args:
            encoder: Instance of a module for encoding
            decoder: Instance of a module for decoding
        """
        super().__init__()
        self.save_hyperparameters()
        self.encoder = encoder
        self.decoder = decoder

If the CLI is implemented as LightningCLI(MyMainModel) the configuration would be as follows:

model:
  encoder:
    class_path: mycode.myencoders.MyEncoder
    init_args:
      ...
  decoder:
    class_path: mycode.mydecoders.MyDecoder
    init_args:
      ...

It is also possible to combine subclass_mode_model=True and submodules, thereby having two levels of class_path.

Tip

By having self.save_hyperparameters() it becomes possible to load the model from a checkpoint. Simply do ModelClass.load_from_checkpoint("path/to/checkpoint.ckpt"). In the case of using subclass_mode_model=True, then load it like LightningModule.load_from_checkpoint("path/to/checkpoint.ckpt"). save_hyperparameters is optional and can be safely removed if there is no need to load from a checkpoint.

Fixed optimizer and scheduler

In some cases, fixing the optimizer and/or learning scheduler might be desired instead of allowing multiple. For this, you can manually add the arguments for specific classes by subclassing the CLI. The following code snippet shows how to implement it:

class MyLightningCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        parser.add_optimizer_args(torch.optim.Adam)
        parser.add_lr_scheduler_args(torch.optim.lr_scheduler.ExponentialLR)

With this, in the config, the optimizer and lr_scheduler groups would accept all of the options for the given classes, in this example, Adam and ExponentialLR. Therefore, the config file would be structured like:

optimizer:
  lr: 0.01
lr_scheduler:
  gamma: 0.2
model:
  ...
trainer:
  ...

where the arguments can be passed directly through the command line without specifying the class. For example:

$ python trainer.py fit --optimizer.lr=0.01 --lr_scheduler.gamma=0.2

Multiple optimizers and schedulers

By default, the CLIs support multiple optimizers and/or learning schedulers, automatically implementing configure_optimizers. This behavior can be disabled by providing auto_configure_optimizers=False on instantiation of LightningCLI. This would be required for example to support multiple optimizers, for each selecting a particular optimizer class. Similar to multiple submodules, this can be done via dependency injection. Unlike the submodules, it is not possible to expect an instance of a class, because optimizers require the module’s parameters to optimize, which are only available after instantiation of the module. Learning schedulers are a similar situation, requiring an optimizer instance. For these cases, dependency injection involves providing a function that instantiates the respective class when called.

An example of a model that uses two optimizers is the following:

from typing import Iterable
from torch.optim import Optimizer


OptimizerCallable = Callable[[Iterable], Optimizer]


class MyModel(LightningModule):
    def __init__(self, optimizer1: OptimizerCallable, optimizer2: OptimizerCallable):
        super().__init__()
        self.save_hyperparameters()
        self.optimizer1 = optimizer1
        self.optimizer2 = optimizer2

    def configure_optimizers(self):
        optimizer1 = self.optimizer1(self.parameters())
        optimizer2 = self.optimizer2(self.parameters())
        return [optimizer1, optimizer2]


cli = MyLightningCLI(MyModel, auto_configure_optimizers=False)

Note the type Callable[[Iterable], Optimizer], which denotes a function that receives a single argument, some learnable parameters, and returns an optimizer instance. With this, from the command line it is possible to select the class and init arguments for each of the optimizers, as follows:

$ python trainer.py fit \
    --model.optimizer1=Adam \
    --model.optimizer1.lr=0.01 \
    --model.optimizer2=AdamW \
    --model.optimizer2.lr=0.0001

In the example above, the OptimizerCallable type alias was created to illustrate what the type hint means. For convenience, this type alias and one for learning schedulers is available in the cli module. An example of a model that uses dependency injection for an optimizer and a learning scheduler is:

from lightning.pytorch.cli import OptimizerCallable, LRSchedulerCallable, LightningCLI


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

    def configure_optimizers(self):
        optimizer = self.optimizer(self.parameters())
        scheduler = self.scheduler(optimizer)
        return {"optimizer": optimizer, "lr_scheduler": scheduler}


cli = MyLightningCLI(MyModel, auto_configure_optimizers=False)

Note that for this example, classes are used as defaults. This is compatible with the type hints, since they are also callables that receive the same first argument and return an instance of the class. Classes that have more than one required argument will not work as default. For these cases a lambda function can be used, e.g. optimizer: OptimizerCallable = lambda p: torch.optim.SGD(p, lr=0.01).

Run from Python

Even though the LightningCLI class is designed to help in the implementation of command line tools, for some use cases it is desired to run directly from Python. To allow this there is the args parameter. An example could be to first implement a normal CLI script, but adding an args parameter with default None to the main function as follows:

from lightning.pytorch.cli import ArgsType, LightningCLI


def cli_main(args: ArgsType = None):
    cli = LightningCLI(MyModel, ..., args=args)
    ...


if __name__ == "__main__":
    cli_main()

Then it is possible to import the cli_main function to run it. Executing in a shell my_cli.py --trainer.max_epochs=100 --model.encoder_layers=24 would be equivalent to:

from my_module.my_cli import cli_main

cli_main(["--trainer.max_epochs=100", "--model.encoder_layers=24"])

All the features that are supported from the command line can be used when giving args as a list of strings. It is also possible to provide a dict or jsonargparse.Namespace. For example in a jupyter notebook someone might do:

args = {
    "trainer": {
        "max_epochs": 100,
    },
    "model": {},
}

args["model"]["encoder_layers"] = 8
cli_main(args)
args["model"]["encoder_layers"] = 12
cli_main(args)
args["trainer"]["max_epochs"] = 200
cli_main(args)

Note

The args parameter must be None when running from command line so that sys.argv is used as arguments. Also, note that the purpose of trainer_defaults is different to args. It is okay to use trainer_defaults in the cli_main function to modify the defaults of some trainer parameters.