The Problem With Raw PyTorch
Training a PyTorch model from scratch means writing hundreds of lines of boilerplate: the training loop, validation loop, optimizer step, gradient zeroing, device management, mixed precision scaler, distributed training setup, checkpoint saving, and metric logging. Every project replicates this code, and every replication introduces subtle bugs.
PyTorch Lightning separates your research code (what your model does) from engineering code (how training runs). Write the model; Lightning handles the rest.
LightningModule: The Core Pattern
import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
class TextClassifier(pl.LightningModule):
def __init__(self, vocab_size: int, hidden_dim: int, num_classes: int, lr: float = 1e-3):
super().__init__()
self.save_hyperparameters()
self.embedding = nn.Embedding(vocab_size, hidden_dim)
self.transformer = nn.TransformerEncoder(
nn.TransformerEncoderLayer(hidden_dim, nhead=8, batch_first=True),
num_layers=4,
)
self.classifier = nn.Linear(hidden_dim, num_classes)
self.loss_fn = nn.CrossEntropyLoss()
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.embedding(x)
x = self.transformer(x)
return self.classifier(x[:, 0]) # CLS token
def training_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = self.loss_fn(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 = self.loss_fn(logits, y)
acc = (logits.argmax(dim=-1) == y).float().mean()
self.log_dict({"val_loss": loss, "val_acc": acc}, prog_bar=True)
def configure_optimizers(self):
optimizer = AdamW(self.parameters(), lr=self.hparams.lr, weight_decay=0.01)
scheduler = CosineAnnealingLR(optimizer, T_max=100)
return {"optimizer": optimizer, "lr_scheduler": scheduler}
