Shortcuts

How to Organize PyTorch Into Lightning

To enable your code to work with Lightning, perform the following to organize PyTorch into Lightning.


1. Move your Computational Code

Move the model architecture and forward pass to your LightningModule.

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F


class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer_1 = nn.Linear(28 * 28, 128)
        self.layer_2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        return x

2. Move the Optimizer(s) and LR Scheduler(s)

Move your optimizers to the configure_optimizers() hook.

class LitModel(pl.LightningModule):
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]

3. Configure the Training Logic

Lightning automates the training loop for you and manages all of the associated components such as: epoch and batch tracking, optimizers and schedulers, and metric reduction. As a user, you just need to define how your model behaves with a batch of training data within the training_step() method. When using Lightning, simply override the training_step() method which takes the current batch and the batch_idx as arguments. Optionally, it can take optimizer_idx if your LightningModule defines multiple optimizers within its configure_optimizers() hook.

class LitModel(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

4. Configure the Validation Logic

Lightning also automates the validation loop for you and manages all of the associated components such as: epoch and batch tracking, and metrics reduction. As a user, you just need to define how your model behaves with a batch of validation data within the validation_step() method. When using Lightning, simply override the validation_step() method which takes the current batch and the batch_idx as arguments. Optionally, it can take dataloader_idx if you configure multiple dataloaders.

To add an (optional) validation loop add logic to the validation_step() hook (make sure to use the hook parameters, batch and batch_idx in this case).

class LitModel(pl.LightningModule):
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        val_loss = F.cross_entropy(y_hat, y)
        self.log("val_loss", val_loss)

Additionally, you can run only the validation loop using validate() method.

model = LitModel()
trainer.validate(model)

Note

model.eval() and torch.no_grad() are called automatically for validation.

Tip

trainer.validate() loads the best checkpoint automatically by default if checkpointing was enabled during fitting.


5. Configure Testing Logic

Lightning automates the testing loop for you and manages all the associated components, such as epoch and batch tracking, metrics reduction. As a user, you just need to define how your model behaves with a batch of testing data within the test_step() method. When using Lightning, simply override the test_step() method which takes the current batch and the batch_idx as arguments. Optionally, it can take dataloader_idx if you configure multiple dataloaders.

class LitModel(pl.LightningModule):
    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        test_loss = F.cross_entropy(y_hat, y)
        self.log("test_loss", test_loss)

The test loop isn’t used within fit(), therefore, you would need to explicitly call test().

model = LitModel()
trainer.test(model)

Note

model.eval() and torch.no_grad() are called automatically for testing.

Tip

trainer.test() loads the best checkpoint automatically by default if checkpointing is enabled.


6. Configure Prediction Logic

Lightning automates the prediction loop for you and manages all of the associated components such as epoch and batch tracking. As a user, you just need to define how your model behaves with a batch of data within the predict_step() method. When using Lightning, simply override the predict_step() method which takes the current batch and the batch_idx as arguments. Optionally, it can take dataloader_idx if you configure multiple dataloaders. If you don’t override predict_step hook, it by default calls forward() method on the batch.

class LitModel(LightningModule):
    def predict_step(self, batch, batch_idx):
        x, y = batch
        pred = self(x)
        return pred

The predict loop will not be used until you call predict().

model = LitModel()
trainer.predict(model)

Note

model.eval() and torch.no_grad() are called automatically for testing.

Tip

trainer.predict() loads the best checkpoint automatically by default if checkpointing is enabled.


7. Remove any .cuda() or .to(device) Calls

Your LightningModule can automatically run on any hardware!

If you have any explicit calls to .cuda() or .to(device), you can remove them since Lightning makes sure that the data coming from DataLoader and all the Module instances initialized inside LightningModule.__init__ are moved to the respective devices automatically. If you still need to access the current device, you can use self.device anywhere in your LightningModule except in the __init__ and setup methods.

class LitModel(LightningModule):
    def training_step(self, batch, batch_idx):
        z = torch.randn(4, 5, device=self.device)
        ...

Hint: If you are initializing a Tensor within the LightningModule.__init__ method and want it to be moved to the device automatically you should call register_buffer() to register it as a parameter.

class LitModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.register_buffer("running_mean", torch.zeros(num_features))

8. Use your own data

To use your DataLoaders, you can override the respective dataloader hooks in the LightningModule:

class LitModel(LightningModule):
    def train_dataloader(self):
        return DataLoader(...)

    def val_dataloader(self):
        return DataLoader(...)

    def test_dataloader(self):
        return DataLoader(...)

    def predict_dataloader(self):
        return DataLoader(...)

Alternatively, you can pass your dataloaders in one of the following ways:

  • Pass in the dataloaders explictly inside trainer.fit/.validate/.test/.predict calls.

  • Use a LightningDataModule.

Checkout Managing Data doc to understand data management within Lightning.

Read the Docs v: stable
Versions
latest
stable
1.6.4
1.6.3
1.6.2
1.6.1
1.6.0
1.5.10
1.5.9
1.5.8
1.5.7
1.5.6
1.5.5
1.5.4
1.5.3
1.5.2
1.5.1
1.5.0
1.4.9
1.4.8
1.4.7
1.4.6
1.4.5
1.4.4
1.4.3
1.4.2
1.4.1
1.4.0
1.3.8
1.3.7
1.3.6
1.3.5
1.3.4
1.3.3
1.3.2
1.3.1
1.3.0
1.2.10
1.2.8
1.2.7
1.2.6
1.2.5
1.2.4
1.2.3
1.2.2
1.2.1
1.2.0
1.1.8
1.1.7
1.1.6
1.1.5
1.1.4
1.1.3
1.1.2
1.1.1
1.1.0
1.0.8
1.0.7
1.0.6
1.0.5
1.0.4
1.0.3
1.0.2
1.0.1
1.0.0
0.10.0
0.9.0
0.8.5
0.8.4
0.8.3
0.8.2
0.8.1
0.8.0
0.7.6
0.7.5
0.7.4
0.7.3
0.7.2
0.7.1
0.7.0
0.6.0
0.5.3
0.4.9
docs_2
Downloads
html
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.