[CLI] Shorthand notation to instantiate callbacks [3/3]

CHANGELOG.md

@@ -37,6 +37,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
     * Added `FastForwardSampler` and `CaptureIterableDataset` injection to data loading utilities ([#8366](https://github.com/PyTorchLightning/pytorch-lightning/pull/8366))
 
 
+- Added simplification for handling optimizers, schedulers and callbacks with `LightningCLI` ([#8815](https://github.com/PyTorchLightning/pytorch-lightning/pull/8815))
+
+
 - Added `CheckpointIO` to expose checkpoint IO from training type plugin ([#8743](https://github.com/PyTorchLightning/pytorch-lightning/pull/8743))
 
 

docs/source/common/lightning_cli.rst

@@ -312,6 +312,42 @@ Similar to the callbacks, any arguments in :class:`~pytorch_lightning.trainer.tr
 :class:`~pytorch_lightning.core.datamodule.LightningDataModule` classes that have as type hint a class can be configured
 the same way using :code:`class_path` and :code:`init_args`.
 
+Lightning optionally simplifies the user command line so that only the :class:`~pytorch_lightning.callbacks.Callback` name is required.
+The argument's order matters and the user needs to pass the arguments in the following way.
+This is supported only for PyTorch Lightning built-in :class:`~pytorch_lightning.callbacks.Callback`.
+
+.. code-block:: bash
+
+    $  python ... --trainer.callbacks={CALLBACK_NAME_1} --trainer.callbacks.{CALLBACK_1_ARGS_1}=... --trainer.callbacks.{CALLBACK_1_ARGS_2}=... --trainer.callbacks={CALLBACK_N} --trainer.callbacks.{CALLBACK_N_ARGS_1}=...
+
+Here is an example:
+
+.. code-block:: bash
+
+    $  python ... --trainer.callbacks=EarlyStopping --trainer.callbacks.patience=5 --trainer.callbacks=LearningRateMonitor
+
+Register your callbacks
+^^^^^^^^^^^^^^^^^^^^^^^
+
+Lightning provides registries for you to add your own callbacks and benefit from the command line simplification as described above:
+
+.. code-block:: python
+
+    from pytorch_lightning.utilities.cli import CALLBACK_REGISTRY
+    from pytorch_lightning.callbacks import Callback
+
+
+    @CALLBACK_REGISTRY
+    class CustomCallback(Callback):
+        ...
+
+
+    cli = LightningCLI(...)
+
+.. code-block:: bash
+
+    $  python ... --trainer.callbacks=CustomCallback ...
+
 
 Multiple models and/or datasets
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -683,6 +719,12 @@ And the same through command line:
 
     $ python train.py --optimizer.class_path=torch.optim.Adam --optimizer.init_args.lr=0.01
 
+Optionally, the command line can be simplified for PyTorch built-in `optimizers` and `schedulers`:
+
+.. code-block:: bash
+
+    $ python train.py --optimizer=Adam --optimizer.lr=0.01
+
 The automatic implementation of :code:`configure_optimizers` can be disabled by linking the configuration group. An
 example can be :code:`ReduceLROnPlateau` which requires to specify a monitor. This would be:
 
@@ -717,13 +759,71 @@ example can be :code:`ReduceLROnPlateau` which requires to specify a monitor. Th
 
     cli = MyLightningCLI(MyModel)
 
+
 For both possibilities of using :meth:`pytorch_lightning.utilities.cli.LightningArgumentParser.add_optimizer_args` with
 a single class or a tuple of classes, the value given to :code:`optimizer_init` will always be a dictionary including
 :code:`class_path` and :code:`init_args` entries. The function
 :func:`~pytorch_lightning.utilities.cli.instantiate_class` takes care of importing the class defined in
 :code:`class_path` and instantiating it using some positional arguments, in this case :code:`self.parameters()`, and the
 :code:`init_args`. Any number of optimizers and learning rate schedulers can be added when using :code:`link_to`.
 
+Built in schedulers & optimizers and registering your own
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+For code simplification, the LightningCLI provides properties with PyTorch's built-in `optimizers` and `schedulers` already registered.
+
+Only the optimizer or scheduler name needs to be passed along its arguments.
+
+.. code-block:: bash
+
+    $ python train.py --optimizer=Adam --optimizer.lr=0.01 --lr_scheduler=CosineAnnealingLR
+
+If your model requires multiple optimizers, the LightningCLI provides already registered optimizers and schedulers under the properties `registered_optimizers` and `registered_lr_schedulers`
+
+.. code-block::
+
+    class MyLightningCLI(LightningCLI):
+        def add_arguments_to_parser(self, parser):
+            parser.add_optimizer_args(
+                self.registered_optimizers,
+                nested="gen_optimizer",
+                link_to="model.optimizer_init",
+            )
+            parser.add_optimizer_args(
+                self.registered_optimizers,
+                nested="gen_discriminator",
+                link_to="model.optimizer_init",
+            )
+
+.. code-block:: bash
+
+    $ python train.py --gen_optimizer=Adam --optimizer.lr=0.01 -gen_discriminator=Adam --optimizer.lr=0.0001
+
+Furthermore, a user can register their own optimizers or schedulers as follows.
+
+.. code-block:: python
+
+    import torch
+    from pytorch_lightning.utilities.cli import OPTIMIZER_REGISTRY, LR_SCHEDULER_REGISTRY
+    from pytorch_lightning.callbacks import Callback
+
+
+    @OPTIMIZER_REGISTRY
+    class CustomAdam(torch.optim.Adam):
+        pass
+
+
+    @LR_SCHEDULER_REGISTRY
+    class CustomCosineAnnealingLR(torch.optim.lr_scheduler.CosineAnnealingLR):
+        pass
+
+
+    cli = LightningCLI(...)
+
+.. code-block:: bash
+
+    $ python train.py --optimizer=CustomAdam --optimizer.lr=0.01 --lr_scheduler=CustomCosineAnnealingLR
+
 
 Notes related to reproducibility
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

pytorch_lightning/plugins/plugins_registry.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import importlib
 import inspect
-from collections import UserDict
 from inspect import getmembers, isclass
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional
@@ -22,7 +21,7 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 
-class _TrainingTypePluginsRegistry(UserDict):
+class _TrainingTypePluginsRegistry(dict):
     """
     This class is a Registry that stores information about the Training Type Plugins.