pytorch-lightning

PyTorch Lightning: Simplified Deep Learning

Overview

PyTorch Lightning provides a lightweight wrapper for PyTorch that handles the training loop complexity. Apply this skill for structured training, reduced boilerplate, automatic device management, built-in logging, and production-ready workflows.

When to Use This Skill

This skill should be used when:

• Reducing training loop boilerplate
• Multi-GPU or TPU training
• Automatic mixed precision training
• Built-in checkpointing and logging
• Structured research code
• Production model deployment
• Reproducible experiments
• Integration with experiment tracking tools

Quick Start

Basic Import and Setup

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

Basic Lightning Module

import pytorch_lightning as pl
import torch
import torch.nn as nn

class LitModel(pl.LightningModule):
    def __init__(self, hidden_dim=256):
        super().__init__()
        self.l1 = nn.Linear(28 * 28, hidden_dim)
        self.l2 = nn.Linear(hidden_dim, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = torch.relu(self.l1(x))
        x = self.l2(x)
        return x

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

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

Training with Trainer

import pytorch_lightning as pl
from torch.utils.data import DataLoader, TensorDataset
from torchvision import transforms
from torchvision.datasets import MNIST

# Data
train_ds = MNIST(".", train=True, transform=transforms.ToTensor(), download=True)
train_loader = DataLoader(train_ds, batch_size=32)

# Model
model = LitModel()

# Trainer
trainer = pl.Trainer(
    max_epochs=5,
    accelerator="auto",  # GPU/CPU auto-detection
)

# Train
trainer.fit(model, train_loader)

LightningModule Structure