Lightning has utilities to interact seamlessly with the command line ArgumentParser and plays well with the hyperparameter optimization framework of your choice.


Lightning is designed to augment a lot of the functionality of the built-in Python ArgumentParser

from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument("--layer_1_dim", type=int, default=128)
args = parser.parse_args()

This allows you to call your program like so:

python --layer_1_dim 64

Argparser Best Practices

It is best practice to layer your arguments in three sections.

  1. Trainer args (gpus, num_nodes, etc…)

  2. Model specific arguments (layer_dim, num_layers, learning_rate, etc…)

  3. Program arguments (data_path, cluster_email, etc…)

We can do this as follows. First, in your LightningModule, define the arguments specific to that module. Remember that data splits or data paths may also be specific to a module (i.e.: if your project has a model that trains on Imagenet and another on CIFAR-10).

class LitModel(LightningModule):
    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("LitModel")
        parser.add_argument("--encoder_layers", type=int, default=12)
        parser.add_argument("--data_path", type=str, default="/some/path")
        return parent_parser

Now in your main trainer file, add the Trainer args, the program args, and add the model args

# ----------------
# ----------------
from argparse import ArgumentParser

parser = ArgumentParser()

# add PROGRAM level args
parser.add_argument("--conda_env", type=str, default="some_name")
parser.add_argument("--notification_email", type=str, default="[email protected]")

# add model specific args
parser = LitModel.add_model_specific_args(parser)

# add all the available trainer options to argparse
# ie: now --gpus --num_nodes ... --fast_dev_run all work in the cli
parser = Trainer.add_argparse_args(parser)

args = parser.parse_args()

Now you can call run your program like so:

python --gpus 2 --num_nodes 2 --conda_env 'my_env' --encoder_layers 12

Finally, make sure to start the training like so:

# init the trainer like this
trainer = Trainer.from_argparse_args(args, early_stopping_callback=...)

# NOT like this
trainer = Trainer(gpus=hparams.gpus, ...)

# init the model with Namespace directly
model = LitModel(args)

# or init the model with all the key-value pairs
dict_args = vars(args)
model = LitModel(**dict_args)

LightningModule hyperparameters

Often times we train many versions of a model. You might share that model or come back to it a few months later at which point it is very useful to know how that model was trained (i.e.: what learning rate, neural network, etc…).

Lightning has a few ways of saving that information for you in checkpoints and yaml files. The goal here is to improve readability and reproducibility.

  1. Using save_hyperparameters() within your LightningModule __init__ function will enable Lightning to store all the provided arguments within the self.hparams attribute. These hyper-parameters will also be stored within the model checkpoint, which simplifies model re-instantiation in production settings. This also makes those values available via self.hparams.

    class LitMNIST(LightningModule):
        def __init__(self, layer_1_dim=128, learning_rate=1e-2, **kwargs):
            # call this to save (layer_1_dim=128, learning_rate=1e-4) to the checkpoint
            # equivalent
            self.save_hyperparameters("layer_1_dim", "learning_rate")
            # Now possible to access layer_1_dim from hparams
  2. Sometimes your init might have objects or other parameters you might not want to save. In that case, choose only a few

    class LitMNIST(LightningModule):
        def __init__(self, loss_fx, generator_network, layer_1_dim=128 ** kwargs):
            self.layer_1_dim = layer_1_dim
            self.loss_fx = loss_fx
            # call this to save (layer_1_dim=128) to the checkpoint
    # to load specify the other args
    model = LitMNIST.load_from_checkpoint(PATH, loss_fx=torch.nn.SomeOtherLoss, generator_network=MyGenerator())
  3. You can also convert full objects such as dict or Namespace to hparams so they get saved to the checkpoint.

    class LitMNIST(LightningModule):
        def __init__(self, conf: Optional[Union[Dict, Namespace, DictConfig]] = None, **kwargs):
            # save the config and any extra arguments
            self.layer_1 = nn.Linear(28 * 28, self.hparams.layer_1_dim)
            self.layer_2 = nn.Linear(self.hparams.layer_1_dim, self.hparams.layer_2_dim)
            self.layer_3 = nn.Linear(self.hparams.layer_2_dim, 10)
    conf = {...}
    # OR
    # conf = parser.parse_args()
    # OR
    # conf = OmegaConf.create(...)
    model = LitMNIST(conf=conf, anything=10)
    # Now possible to access any stored variables from hparams
    # for this to work, you need to access with `self.hparams.layer_1_dim`, not `conf.layer_1_dim`
    model = LitMNIST.load_from_checkpoint(PATH)

Trainer args

To recap, add ALL possible trainer flags to the argparser and init the Trainer this way

parser = ArgumentParser()
parser = Trainer.add_argparse_args(parser)
hparams = parser.parse_args()

trainer = Trainer.from_argparse_args(hparams)

# or if you need to pass in callbacks
trainer = Trainer.from_argparse_args(hparams, enable_checkpointing=..., callbacks=[...])

Multiple Lightning Modules

We often have multiple Lightning Modules where each one has different arguments. Instead of polluting the file, the LightningModule lets you define arguments for each one.

class LitMNIST(LightningModule):
    def __init__(self, layer_1_dim, **kwargs):
        self.layer_1 = nn.Linear(28 * 28, layer_1_dim)

    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("LitMNIST")
        parser.add_argument("--layer_1_dim", type=int, default=128)
        return parent_parser
class GoodGAN(LightningModule):
    def __init__(self, encoder_layers, **kwargs):
        self.encoder = Encoder(layers=encoder_layers)

    def add_model_specific_args(parent_parser):
        parser = parent_parser.add_argument_group("GoodGAN")
        parser.add_argument("--encoder_layers", type=int, default=12)
        return parent_parser

Now we can allow each model to inject the arguments it needs in the

def main(args):
    dict_args = vars(args)

    # pick model
    if args.model_name == "gan":
        model = GoodGAN(**dict_args)
    elif args.model_name == "mnist":
        model = LitMNIST(**dict_args)

    trainer = Trainer.from_argparse_args(args)

if __name__ == "__main__":
    parser = ArgumentParser()
    parser = Trainer.add_argparse_args(parser)

    # figure out which model to use
    parser.add_argument("--model_name", type=str, default="gan", help="gan or mnist")

    temp_args, _ = parser.parse_known_args()

    # let the model add what it wants
    if temp_args.model_name == "gan":
        parser = GoodGAN.add_model_specific_args(parser)
    elif temp_args.model_name == "mnist":
        parser = LitMNIST.add_model_specific_args(parser)

    args = parser.parse_args()

    # train

and now we can train MNIST or the GAN using the command line interface!

$ python --model_name gan --encoder_layers 24
$ python --model_name mnist --layer_1_dim 128