mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-09 11:32:07 +08:00
William Falcon
@@ -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
|
||||
^^^^^^^^^^^^
|
||||
|
||||
@@ -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 <hyperparameters.rst>`_.
|
||||
.. include:: hyperparameters.rst
|
||||
|
||||
---------
|
||||
|
||||
|
||||
+242
-47
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user