Changed basic_examples to use LightningCLI (#6862)

Co-authored-by: Carlos Mocholi <carlossmocholi@gmail.com>

azure-pipelines.yml

@@ -116,8 +116,8 @@ jobs:
  set -e
  python -m pytest pl_examples -v --maxfail=2 --durations=0
  pip install . --user --quiet
- bash pl_examples/run_examples-args.sh --gpus 1 --max_epochs 1 --batch_size 64 --limit_train_batches 5 --limit_val_batches 3
- bash pl_examples/run_ddp-examples.sh --max_epochs 1 --batch_size 32 --limit_train_batches 2 --limit_val_batches 2
+ bash pl_examples/run_examples-args.sh --trainer.gpus 1 --trainer.max_epochs 1 --data.batch_size 64 --trainer.limit_train_batches 5 --trainer.limit_val_batches 3
+ bash pl_examples/run_ddp-examples.sh --trainer.max_epochs 1 --data.batch_size 32 --trainer.limit_train_batches 2 --trainer.limit_val_batches 2
  # cd pl_examples/basic_examples
  # bash submit_ddp_job.sh
  # bash submit_ddp2_job.sh

docs/source/common/lightning_cli.rst

@@ -224,7 +224,7 @@ datamodule class. However, there are many cases in which the objective is to eas
 multiple models and datasets. For these cases the tool can be configured such that a model and/or a datamodule is
 specified by an import path and init arguments. For example, with a tool implemented as:
 
-.. testcode::
+.. code-block:: python
 
  from pytorch_lightning.utilities.cli import LightningCLI
 

pl_examples/basic_examples/README.md

@@ -8,10 +8,10 @@ Trains MNIST where the model is defined inside the `LightningModule`.
 python simple_image_classifier.py
 
 # gpus (any number)
-python simple_image_classifier.py --gpus 2
+python simple_image_classifier.py --trainer.gpus 2
 
 # dataparallel
-python simple_image_classifier.py --gpus 2 --distributed_backend 'dp'
+python simple_image_classifier.py --trainer.gpus 2 --trainer.accelerator 'dp'
 ```
 
 ---
@@ -30,10 +30,10 @@ Generic image classifier with an arbitrary backbone (ie: a simple system)
 python backbone_image_classifier.py
 
 # gpus (any number)
-python backbone_image_classifier.py --gpus 2
+python backbone_image_classifier.py --trainer.gpus 2
 
 # dataparallel
-python backbone_image_classifier.py --gpus 2 --distributed_backend 'dp'
+python backbone_image_classifier.py --trainer.gpus 2 --trainer.accelerator 'dp'
 ```
 
 ---
@@ -44,10 +44,10 @@ Showing the power of a system... arbitrarily complex training loops
 python autoencoder.py
 
 # gpus (any number)
-python autoencoder.py --gpus 2
+python autoencoder.py --trainer.gpus 2
 
 # dataparallel
-python autoencoder.py --gpus 2 --distributed_backend 'dp'
+python autoencoder.py --trainer.gpus 2 --trainer.accelerator 'dp'
 ```
 ---
 # Multi-node example

pl_examples/basic_examples/autoencoder.py

@@ -11,8 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""
+MNIST autoencoder example.
 
-from argparse import ArgumentParser
+To run:
+python autoencoder.py --trainer.max_epochs=50
+"""
 
 import torch
 import torch.nn.functional as F
@@ -21,6 +25,7 @@
 
 import pytorch_lightning as pl
 from pl_examples import _DATASETS_PATH, _TORCHVISION_MNIST_AVAILABLE, cli_lightning_logo
+from pytorch_lightning.utilities.cli import LightningCLI
 from pytorch_lightning.utilities.imports import _TORCHVISION_AVAILABLE
 
 if _TORCHVISION_AVAILABLE:
@@ -87,44 +92,31 @@ def configure_optimizers(self):
  return optimizer
 
 
+class MyDataModule(pl.LightningDataModule):
+
+ def __init__(
+ self,
+ batch_size: int = 32,
+ ):
+ super().__init__()
+ dataset = MNIST(_DATASETS_PATH, train=True, download=True, transform=transforms.ToTensor())
+ self.mnist_test = MNIST(_DATASETS_PATH, train=False, download=True, transform=transforms.ToTensor())
+ self.mnist_train, self.mnist_val = random_split(dataset, [55000, 5000])
+ self.batch_size = batch_size
+
+ def train_dataloader(self):
+ return DataLoader(self.mnist_train, batch_size=self.batch_size)
+
+ def val_dataloader(self):
+ return DataLoader(self.mnist_val, batch_size=self.batch_size)
+
+ def test_dataloader(self):
+ return DataLoader(self.mnist_test, batch_size=self.batch_size)
+
+
 def cli_main():
- pl.seed_everything(1234)
-
- # ------------
- # args
- # ------------
- parser = ArgumentParser()
- parser.add_argument('--batch_size', default=32, type=int)
- parser.add_argument('--hidden_dim', type=int, default=64)
- parser = pl.Trainer.add_argparse_args(parser)
- args = parser.parse_args()
-
- # ------------
- # data
- # ------------
- dataset = MNIST(_DATASETS_PATH, train=True, download=True, transform=transforms.ToTensor())
- mnist_test = MNIST(_DATASETS_PATH, train=False, download=True, transform=transforms.ToTensor())
- mnist_train, mnist_val = random_split(dataset, [55000, 5000])
-
- train_loader = DataLoader(mnist_train, batch_size=args.batch_size)
- val_loader = DataLoader(mnist_val, batch_size=args.batch_size)
- test_loader = DataLoader(mnist_test, batch_size=args.batch_size)
-
- # ------------
- # model
- # ------------
- model = LitAutoEncoder(args.hidden_dim)
-
- # ------------
- # training
- # ------------
- trainer = pl.Trainer.from_argparse_args(args)
- trainer.fit(model, train_loader, val_loader)
-
- # ------------
- # testing
- # ------------
- result = trainer.test(test_dataloaders=test_loader)
+ cli = LightningCLI(LitAutoEncoder, MyDataModule, seed_everything_default=1234)
+ result = cli.trainer.test(cli.model, datamodule=cli.datamodule)
  print(result)
 
 

pl_examples/basic_examples/backbone_image_classifier.py

@@ -11,15 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""
+MNIST backbone image classifier example.
 
-from argparse import ArgumentParser
+To run:
+python backbone_image_classifier.py --trainer.max_epochs=50
+"""
 
 import torch
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, random_split
 
 import pytorch_lightning as pl
 from pl_examples import _DATASETS_PATH, _TORCHVISION_MNIST_AVAILABLE, cli_lightning_logo
+from pytorch_lightning.utilities.cli import LightningCLI
 from pytorch_lightning.utilities.imports import _TORCHVISION_AVAILABLE
 
 if _TORCHVISION_AVAILABLE:
@@ -59,7 +64,11 @@ class LitClassifier(pl.LightningModule):
  )
  """
 
- def __init__(self, backbone, learning_rate=1e-3):
+ def __init__(
+ self,
+ backbone,
+ learning_rate: float = 0.0001,
+ ):
  super().__init__()
  self.save_hyperparameters()
  self.backbone = backbone
@@ -92,52 +101,42 @@ def configure_optimizers(self):
  # self.hparams available because we called self.save_hyperparameters()
  return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
 
- @staticmethod
- def add_model_specific_args(parent_parser):
- parser = parent_parser.add_argument_group("LitClassifier")
- parser.add_argument('--learning_rate', type=float, default=0.0001)
- return parent_parser
+
+class MyDataModule(pl.LightningDataModule):
+
+ def __init__(
+ self,
+ batch_size: int = 32,
+ ):
+ super().__init__()
+ dataset = MNIST(_DATASETS_PATH, train=True, download=True, transform=transforms.ToTensor())
+ self.mnist_test = MNIST(_DATASETS_PATH, train=False, download=True, transform=transforms.ToTensor())
+ self.mnist_train, self.mnist_val = random_split(dataset, [55000, 5000])
+ self.batch_size = batch_size
+
+ def train_dataloader(self):
+ return DataLoader(self.mnist_train, batch_size=self.batch_size)
+
+ def val_dataloader(self):
+ return DataLoader(self.mnist_val, batch_size=self.batch_size)
+
+ def test_dataloader(self):
+ return DataLoader(self.mnist_test, batch_size=self.batch_size)
+
+
+class MyLightningCLI(LightningCLI):
+
+ def add_arguments_to_parser(self, parser):
+ parser.add_class_arguments(Backbone, 'model.backbone')
+
+ def instantiate_model(self):
+ self.config_init['model']['backbone'] = Backbone(**self.config['model']['backbone'])
+ super().instantiate_model()
 
 
 def cli_main():
- pl.seed_everything(1234)
-
- # ------------
- # args
- # ------------
- parser = ArgumentParser()
- parser.add_argument('--batch_size', default=32, type=int)
- parser.add_argument('--hidden_dim', type=int, default=128)
- parser = pl.Trainer.add_argparse_args(parser)
- parser = LitClassifier.add_model_specific_args(parser)
- args = parser.parse_args()
-
- # ------------
- # data
- # ------------
- dataset = MNIST(_DATASETS_PATH, train=True, download=True, transform=transforms.ToTensor())
- mnist_test = MNIST(_DATASETS_PATH, train=False, download=True, transform=transforms.ToTensor())
- mnist_train, mnist_val = random_split(dataset, [55000, 5000])
-
- train_loader = DataLoader(mnist_train, batch_size=args.batch_size)
- val_loader = DataLoader(mnist_val, batch_size=args.batch_size)
- test_loader = DataLoader(mnist_test, batch_size=args.batch_size)
-
- # ------------
- # model
- # ------------
- model = LitClassifier(Backbone(hidden_dim=args.hidden_dim), args.learning_rate)
-
- # ------------
- # training
- # ------------
- trainer = pl.Trainer.from_argparse_args(args)
- trainer.fit(model, train_loader, val_loader)
-
- # ------------
- # testing
- # ------------
- result = trainer.test(test_dataloaders=test_loader)
+ cli = MyLightningCLI(LitClassifier, MyDataModule, seed_everything_default=1234)
+ result = cli.trainer.test(cli.model, datamodule=cli.datamodule)
  print(result)