PyTorch Lightning
Gregor Lenz
26.7.2022
PyTorch Lightning
- Extends the 'low-level' PyTorch functionality
- Standardises neural network train / test setup
- Provides many helpful built-in functions
- Reduces boilerplate code
Who might use it
- Anyone who has previously trained neural networks with pure PyTorch
model = nn.Sequential(nn.Linear(...), nn.Linear(...))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
device = 'cuda:0'
for epoch in range(10):
model.train()
for data, y_true in iter(training_dataloader):
data = data.to(device)
targets = targets.to(device)
optimizer.zero_grad()
y_hat = model(data)
loss = criterion(y_hat, y_true)
loss.backward()
optimizer.step()
# log the loss
...
model.eval()
# log the epoch loss, evaluate
...Problem: many ways to do the same thing
- Logging parameters, scalars
- Saving models
- Moving data between devices
- Training on multiple devices
- Repetitive code for training, validation, testing
Lightning makes a few assumptions
-
There is at least a training loop of data that we cycle through, potentially also validation and testing
-
If we specify a device, we are going to do all the computations on there
-
Statistics are accumulated per epoch
-
For every iteration:
- gradients are set to zero
- loss is backpropagated
- optimizer takes a step
import pytorch_lightning as pl
class MyLightningModel(pl.LightningModule):
def __init__(self):
# define model here...
def forward(self, data: torch.Tensor):
# specify how data is fed to the model
def training_step(self, batch):
# unpack a mini-batch, feed data, calculate losses
def validation_step(self, batch):
# unpack a mini-batch, feed data, calculate losses,
# calculate prediction accuracy
def configure_optimizers(self):
# specify optimizer here