Extend the LightningCLI to register models, datamodules, and callbacks.

[tchaton]

🚀 Feature

Implement a CLI store/provider/registry for available LightningModules, LightningDataModules, and Callbacks.

Implementation

class LightningCLI:
    store = Registry()

    @classmethod
    def register(cls, *args, **kwargs):
        cls.store.register(*args, **kwargs)

The registry idea is highly influenced by this implementation:

https://github.com/PyTorchLightning/lightning-flash/blob/master/flash/core/registry.py
https://github.com/PyTorchLightning/lightning-flash/blob/master/tests/core/test_registry.py

Usage

# note: these `register` calls could be in different files
LightningCLI.register('my-callback', CallbackA)

LightningCLI.register('vision-model', VisionModel)

LightningCLI.register('mnist-data', MnistDataModule)
LightningCLI.register('cifar10-data', CIFAR10DataModule)
LightningCLI.register('imagenet-data', ImageNetDataModule)

# API might change - naming the model-data pair as an experiment
LightningCLI.register_experiment(VisionModel, MnistDataModule)
LightningCLI.register_experiment(VisionModel, CIFAR10DataModule)
LightningCLI.register_experiment(VisionModel, ImageneteDataModule)

# or with a decorator.
# `name` could be taken from the class or function
# `arg` could be inferred from the parent class
@LightningCLI.register(name='some_model', arg='model')
class SomeModel(LightningModule)
    ...

# in train.py
cli = LightningCLI()

Console interaction

Basic help

python train.py -h 

callbacks: {
    'my-callback': CallbackA,
}
models: {
    'vision-model': VisionModel,
}
datamodules: {
    'mnist-data': MnistDataModule,
    'cifar10-data': CIFAR10DataModule,
    'imagenet-data': ImageNetDataModule,
}
experiments: {
    'vision-mnist': {'model': VisionModel, 'data': MnistDataModule},
    'vision-cifar10': {'model': VisionModel, 'data': CIFAR10DataModule},
    'vision-imagenet': {'model': VisionModel, 'data': ImageNetDataModule},  
}

Note that the help output should look nicer, showing raw dicts for simplicity

Filter by category

python train.py -h model=VisionModel
# alternatively
python train-py -h model='vision-model'

callbacks: {
    'my-callback': CallbackA,
}
datamodules: {
    'mnist-data': MnistDataModule,
    'cifar10-data': CIFAR10DataModule,
    'imagenet-data': ImageNetDataModule,
}
experiments: {
    'vision-mnist': {'model': VisionModel, 'data': MnistDataModule},
    'vision-cifar10': {'model': VisionModel, 'data': CIFAR10DataModule},
    'vision-imagenet': {'model': VisionModel, 'data': ImageNetDataModule},  
}

python train.py -h data=CIFAR10DataModule

callbacks: {
    'my-callback': CallbackA,
}
models: {
    'vision-model': VisionModel,
}
datamodules: {
    'cifar10-data': CIFAR10DataModule,
}
experiments: {
    'vision-cifar10': {'model': VisionModel, 'data': CIFAR10DataModule},
}

Basic LightningCLI usage (similar to what's currently implemented)

python train.py -h model=VisionModel data=ImageNetDataModule
# same as
python train.py -h experiment='vision-imagenet'

[mauvilsa]

@tchaton thank you for the proposal. I haven't thought about it yet. But I will do and comment later.

[carmocca]

@tchaton @mauvilsa Modified the original post

[mpariente]

This looks very nice, Asteroid would most probably use this!

[Andry-Bal]

This feature also will make specifying arguments in command line more convenient. Currently if you use subclasses for either model or data, for every argument you have to use init_args in command line, e.g. --data.init_args.batch_size 64, which is a bit annoying.

[mauvilsa]

Being able to register models or data modules does not imply that the arguments can be overridden without using init_args. That is completely different feature. If you like you can create a feature request for this. Could be in jsonargparse since it does not really relate to lightning.

[tchaton]

@register_model('image')
ModelA

@register_model('image')
ModelB

@register_datamodule('image')
DataModuleA

@register_datamodule('image')
DataModuleB

model = register.create_model()
datamodule = register.create_datamodule()
callbacks = register.create_callbacks()
trainer = register.create_trainer()