From d2b94ca81bf75bd052724e66f1ec585388800f08 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sun, 26 Apr 2020 10:57:26 -0400 Subject: [PATCH] clean up docs (#1614) * fixed hparams section * docs clean up --- README.md | 5 +- docs/source/hyperparameters.rst | 18 +- docs/source/introduction_guide.rst | 66 +------ docs/source/new-project.rst | 289 ++++++++++++++++++++++++----- 4 files changed, 254 insertions(+), 124 deletions(-) diff --git a/README.md b/README.md index 2de51a2f..b066e797 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,8 @@ pip install pytorch-lightning [MNIST on TPUs](https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3) ## What is it? -Lightning is a way to organize your PyTorch code to decouple the science code from the engineering. It's more of a style-guide than a framework. +Lightning is a way to organize your PyTorch code to decouple the science code from the engineering. +It's more of a PyTorch style-guide than a framework. In Lightning, you organize your code into 3 distinct categories: @@ -69,6 +70,8 @@ Here's an example of how to refactor your research code into a [LightningModule] The rest of the code is automated by the [Trainer](https://pytorch-lightning.readthedocs.io/en/latest/trainer.html)! ![PT to PL](docs/source/_images/lightning_module/pt_trainer.png) +[READ THIS QUICK START PAGE](https://pytorch-lightning.readthedocs.io/en/latest/new-project.html) + ## Testing Rigour All the automated code by the Trainer is [tested rigorously with every new PR](https://github.com/PyTorchLightning/pytorch-lightning/tree/master/tests). diff --git a/docs/source/hyperparameters.rst b/docs/source/hyperparameters.rst index 7a20b5f6..ff067b82 100644 --- a/docs/source/hyperparameters.rst +++ b/docs/source/hyperparameters.rst @@ -27,9 +27,9 @@ 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...) +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 @@ -84,15 +84,11 @@ Finally, make sure to start the training like so: # YES model = LitModel(hparams) - - # NO - # model = LitModel(learning_rate=hparams.learning_rate, ...) - - # YES trainer = Trainer.from_argparse_args(hparams, early_stopping_callback=...) # NO - trainer = Trainer(gpus=hparams.gpus, ...) + # model = LitModel(learning_rate=hparams.learning_rate, ...) + #trainer = Trainer(gpus=hparams.gpus, ...) LightiningModule hparams @@ -144,8 +140,8 @@ Now pass in the params when you init your model The line `self.hparams = hparams` is very special. This line assigns your hparams to the LightningModule. This does two things: - 1. It adds them automatically to tensorboard logs under the hparams tab. - 2. Lightning will save those hparams to the checkpoint and use them to restore the module correctly. +1. It adds them automatically to tensorboard logs under the hparams tab. +2. Lightning will save those hparams to the checkpoint and use them to restore the module correctly. Trainer args ^^^^^^^^^^^^ diff --git a/docs/source/introduction_guide.rst b/docs/source/introduction_guide.rst index a98631ec..bed1f477 100644 --- a/docs/source/introduction_guide.rst +++ b/docs/source/introduction_guide.rst @@ -604,71 +604,7 @@ Notice the epoch is MUCH faster! --------- -Hyperparameters ---------------- -Normally, we don't hard-code the values to a model. We usually use the command line to -modify the network. - -.. code-block:: python - - from argparse import ArgumentParser - - parser = ArgumentParser() - - # parametrize the network - parser.add_argument('--layer_1_dim', type=int, default=128) - parser.add_argument('--layer_2_dim', type=int, default=256) - parser.add_argument('--batch_size', type=int, default=64) - - args = parser.parse_args() - -Now we can parametrize the LightningModule. - -.. code-block:: python - :emphasize-lines: 5,6,7,12,14 - - class LitMNIST(pl.LightningModule): - def __init__(self, hparams): - super().__init__() - self.hparams = hparams - - self.layer_1 = torch.nn.Linear(28 * 28, hparams.layer_1_dim) - self.layer_2 = torch.nn.Linear(hparams.layer_1_dim, hparams.layer_2_dim) - self.layer_3 = torch.nn.Linear(hparams.layer_2_dim, 10) - - def forward(self, x): - ... - - def train_dataloader(self): - ... - return DataLoader(mnist_train, batch_size=self.hparams.batch_size) - - def configure_optimizers(self): - return Adam(self.parameters(), lr=self.hparams.learning_rate) - - hparams = parse_args() - model = LitMNIST(hparams) - -.. note:: Bonus! if (hparams) is in your module, Lightning will save it into the checkpoint and restore your - model using those hparams exactly. - -And we can also add all the flags available in the Trainer to the Argparser. - -.. code-block:: python - - # add all the available Trainer options to the ArgParser - parser = pl.Trainer.add_argparse_args(parser) - args = parser.parse_args() - -And now you can start your program with - -.. code-block:: bash - - # now you can use any trainer flag - $ python main.py --num_nodes 2 --gpus 8 - - -For a full guide on using hyperparameters, `check out the hyperparameters docs `_. +.. include:: hyperparameters.rst --------- diff --git a/docs/source/new-project.rst b/docs/source/new-project.rst index c269c819..ab066a0b 100644 --- a/docs/source/new-project.rst +++ b/docs/source/new-project.rst @@ -1,72 +1,267 @@ Quick Start =========== -| To start a new project define two files, a LightningModule and a Trainer file. -| To illustrate the power of Lightning and its simplicity, here's an example of a typical research flow. -Case 1: BERT ------------- +PyTorch Lightning is nothing more than organized PyTorch code. +Once you've organized it into a LightningModule, it automates most of the training for you. -| Let's say you're working on something like BERT but want to try different ways of training or even different networks. -| You would define a single LightningModule and use flags to switch between your different ideas. +To illustrate, here's the typical PyTorch project structure organized in a LightningModule. + +.. figure:: /_images/mnist_imgs/pt_to_pl.jpg + :alt: Convert from PyTorch to Lightning + + +Step 1: Define a LightningModule +--------------------------------- .. code-block:: python - class BERT(pl.LightningModule): - def __init__(self, model_name, task): - self.task = task + import os - if model_name == 'transformer': - self.net = Transformer() - elif model_name == 'my_cool_version': - self.net = MyCoolVersion() + import torch + from torch.nn import functional as F + from torch.utils.data import DataLoader + from torchvision.datasets import MNIST + from torchvision import transforms - def training_step(self, batch, batch_idx): - if self.task == 'standard_bert': - # do standard bert training with self.net... - # return loss + import pytorch_lightning as pl - if self.task == 'my_cool_task': - # do my own version with self.net - # return loss + class LitModel(pl.LightningModule): - -Case 2: COOLER NOT BERT ------------------------ - -But if you wanted to try something **completely** different, you'd define a new module for that. - - -.. code-block:: python - - class CoolerNotBERT(pl.LightningModule): def __init__(self): - self.net = ... + super().__init__() + self.l1 = torch.nn.Linear(28 * 28, 10) + + def forward(self, x): + return torch.relu(self.l1(x.view(x.size(0), -1))) def training_step(self, batch, batch_idx): - # do some other cool task - # return loss + x, y = batch + y_hat = self(x) + loss = F.cross_entropy(y_hat, y) + tensorboard_logs = {'train_loss': loss} + return {'loss': loss, 'log': tensorboard_logs} + + def configure_optimizers(self): + return torch.optim.Adam(self.parameters(), lr=0.001) + + def train_dataloader(self): + dataset = MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()) + loader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True) + return loader -Rapid research flow -------------------- - -Then you could do rapid research by switching between these two and using the same trainer. - +Step 2: Fit with a Trainer +-------------------------- .. code-block:: python - if use_bert: - model = BERT() - else: - model = CoolerNotBERT() + from pytorch_lightning import Trainer - trainer = Trainer(gpus=4, precision=16) + model = LitModel() + + # most basic trainer, uses good defaults + trainer = Trainer(gpus=8, num_nodes=1) trainer.fit(model) + Under the hood, lightning does (in high-level pseudocode): -**Notice a few things about this flow:** +.. code-block:: python -1. You're writing pure PyTorch... no unnecessary abstractions or new libraries to learn. -2. You get free GPU and 16-bit support without writing any of that code in your model. -3. You also get early stopping, multi-gpu training, 16-bit and MUCH more without coding anything! + model = LitModel() + train_dataloader = model.train_dataloader + optimizer = model.configure_optimizers() + for epoch in epochs: + train_outs = [] + for batch in train_dataloader: + loss = model.training_step() + loss.backward() + train_outs.append(loss.detach()) + + optimizer.step() + optimizer.zero_grad() + + # optional for logging, etc... + model.training_epoch_end(train_outs) + +Validation loop +--------------- +To also add a validation loop add the following functions + +.. code-block:: python + + class LitModel(pl.LightningModule): + + def validation_step(self, batch, batch_idx): + x, y = batch + y_hat = self(x) + return {'val_loss': F.cross_entropy(y_hat, y)} + + def validation_epoch_end(self, outputs): + avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean() + tensorboard_logs = {'val_loss': avg_loss} + return {'avg_val_loss': avg_loss, 'log': tensorboard_logs + + def val_dataloader(self): + # TODO: do a real train/val split + dataset = MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()) + loader = DataLoader(dataset, batch_size=32, num_workers=4) + return loader + +And now the trainer will call the validation loop automatically + +.. code-block:: python + + # most basic trainer, uses good defaults + trainer = Trainer(gpus=8, num_nodes=1) + trainer.fit(model) + +Under the hood in pseudocode, lightning does the following: + +.. code-block:: python + + # ... + for batch in train_dataloader: + loss = model.training_step() + loss.backward() + # ... + + if validate_at_some_point: + model.eval() + val_outs = [] + for val_batch in model.val_dataloader: + val_out = model.validation_step(val_batch) + val_outs.append(val_out) + + model.validation_epoch_end(val_outs) + model.train() + +The beauty of Lightning is that it handles the details of when to validate, when to call .eval(), +turning off gradients, detaching graphs, making sure you don't enable shuffle for val, etc... + +ie: Lightning removes all the million details you need to remember during research + +Test loop +--------- +You might also need a test loop + +.. code-block:: python + + class LitModel(pl.LightningModule): + + def test_step(self, batch, batch_idx): + x, y = batch + y_hat = self(x) + return {'test_loss': F.cross_entropy(y_hat, y)} + + def test_epoch_end(self, outputs): + avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean() + tensorboard_logs = {'test_loss': avg_loss} + return {'avg_test_loss': avg_loss, 'log': tensorboard_logs} + + def test_dataloader(self): + # TODO: do a real train/val split + dataset = MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()) + loader = DataLoader(dataset, batch_size=32, num_workers=4) + return loader + +However, this time you need to specifically call test (this is done so you don't use the test set by mistake) + +.. code-block:: python + + # OPTION 1: + # test after fit + trainer.fit(model) + trainer.test() + + # OPTION 2: + # test after loading weights + model = LitModel.load_from_checkpoint(PATH) + trainer = Trainer(num_tpu_cores=1) + trainer.test() + +Again, under the hood, lightning does the following in (pseudocode): + +.. code-block:: python + + model.eval() + test_outs = [] + for test_batch in model.test_dataloader: + test_out = model.test_step(val_batch) + test_outs.append(test_out) + + model.test_epoch_end(test_outs) + +Datasets +-------- +If you don't want to define the datasets as part of the LightningModule, just pass them into fit instead. + +.. code-block:: python + + # pass in datasets if you want. + train_dataloader = DataLoader(dataset, batch_size=32, num_workers=4) + val_dataloader, test_dataloader = ... + + trainer = Trainer(gpus=8, num_nodes=1) + trainer.fit(model, train_dataloader, val_dataloader) + + trainer.test(test_dataloader=test_dataloader) + +The advantage of this method is the ability to reuse models for different datasets. The disadvantage +is that for research it makes readability and reproducibility more difficult. This is why we recommend +to define the datasets in the LightningModule if you're doing research, but use the method above for +production models or for prediction tasks. + +Why do you need Lightning? +-------------------------- +Notice the code above has nothing about .cuda() or 16-bit or early stopping or logging, etc... +This is where Lightning adds a ton of value. + +Without changing a SINGLE line of your code, you can now do the following with the above code + +.. code-block:: python + + # train on TPUs using 16 bit precision with early stopping + # using only half the training data and checking validation every quarter of a training epoch + trainer = Trainer( + nb_tpu_cores=8, + precision=16, + early_stop_checkpoint=True, + train_percent_check=0.5, + val_check_interval=0.25 + ) + + # train on 256 GPUs + trainer = Trainer( + gpus=8, + num_nodes=32 + ) + + # train on 1024 CPUs across 128 machines + trainer = Trainer( + num_processes=8, + num_nodes=128 + ) + +And the best part is that your code is STILL just PyTorch... meaning you can do anything you +would normally do. + +.. code-block:: python + + model = LitModel() + model.eval() + + y_hat = model(x) + + model.anything_you_can_do_with_pytorch() + +Summary +------- +In short, by refactoring your PyTorch code: + +1. You STILL keep pure PyTorch. +2. You DON't lose any flexibility. +3. You can get rid of all of your boilerplate. +4. You make your code generalizable to any hardware. +5. Your code is now readable and easier to reproduce (ie: you help with the reproducibility crisis). +6. Your LightningModule is still just a pure PyTorch module.