PyTorch Lightning
PyTorch Lightning factors the boilerplate (training loop, distributed setup, checkpointing, logging) out of your research code so your LightningModule only describes what your model does. The same code runs on CPU, single GPU, multi-GPU, or TPU with no rewrite — Lightning handles the device placement.
A minimal LightningModule and Trainer
EXAMPLE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import lightning as L
class MLP(L.LightningModule):
def __init__(self, in_dim=10, hidden=32, out_dim=2, lr=1e-3):
super().__init__()
self.save_hyperparameters()
self.net = nn.Sequential(
nn.Linear(in_dim, hidden), nn.ReLU(),
nn.Linear(hidden, out_dim),
)
def forward(self, x):
return self.net(x)
def training_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.cross_entropy(logits, y)
self.log('train_loss', loss, prog_bar=True)
return loss
def validation_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.cross_entropy(logits, y)
acc = (logits.argmax(-1) == y).float().mean()
self.log_dict({'val_loss': loss, 'val_acc': acc}, prog_bar=True)
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters(), lr=self.hparams.lr)
# Fake data
X = torch.randn(1024, 10); y = torch.randint(0, 2, (1024,))
train = DataLoader(TensorDataset(X[:800], y[:800]), batch_size=32)
val = DataLoader(TensorDataset(X[800:], y[800:]), batch_size=32)
model = MLP()
trainer = L.Trainer(max_epochs=5, accelerator='auto', devices='auto')
trainer.fit(model, train, val)
Why it matters
save_hyperparameters() is the unsung hero — it captures init args into self.hparams and serialises them with the checkpoint, so you can later restore an identical model from the .ckpt file without retyping the config.
Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.
Example
Example
import lightning as L
class Lit(L.LightningModule):
def __init__(self): super().__init__(); self.net = MLP()
def training_step(self, batch, _):
x, y = batch; return F.cross_entropy(self.net(x), y)
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters(), lr=1e-3)
L.Trainer(max_epochs=10).fit(Lit(), train_loader)
Try it Yourself »
Discussion
Loading…