02 · The LightningModule 🧩

5 min read

This is where your model lives in Lightning. The trick to learning it: every method maps directly onto a piece of the loop you already write by hand.

1. The 'Why'

The LightningModule is the heart of Lightning, and the fastest way to understand it is to realize it's not a new concept — it's your existing training code, cut along its natural seams and dropped into named methods. When you write a raw loop, you already mentally distinguish "the part that computes a training step's loss" from "the part that sets up the optimizer" from "the part that runs a validation step." Lightning simply asks you to make those divisions explicit by putting each into its own method. Then, instead of you calling them in the right order inside hand-written loops, the Trainer (sub-module 03) calls them for you — at the right time, on the right device, with zero_grad/backward/step handled automatically.

The payoff is that your model code becomes a clean, declarative description of what your model does with none of the loop mechanics tangled in. A LightningModule subclass holds your nn.Module layers (in __init__, exactly as before), defines a training_step that returns the loss for one batch, optionally a validation_step for metrics, and a configure_optimizers that returns your optimizer. Notice what's missing: no for epoch, no .to(device), no zero_grad, no backward, no eval(). Lightning supplies all of that. This sub-module shows the one-to-one mapping from the raw loop you mastered in Module 05 to these methods — once you see the correspondence, the LightningModule stops looking like framework magic and starts looking like tidy PyTorch.

2. Core Concepts

It is an nn.Module. LightningModule subclasses nn.Module, so everything from Module 03 applies — you define layers in __init__ and can still call the model directly. It just adds structured hooks.

training_step(self, batch, batch_idx). The body of one training iteration: unpack the batch, run the forward pass, compute the loss, and return the loss. Lightning does zero_grad, backward, and step around it. The batch is already on the correct device.

validation_step(self, batch, batch_idx). Like training_step but for validation — compute and log metrics. Lightning wraps it in eval() + no_grad() automatically. (test_step mirrors it for the test set.)

configure_optimizers(self). Create and return your optimizer(s) (and optionally LR schedulers). Replaces the manual optimizer setup from Module 05.

self.log(...). Report a metric by name; Lightning aggregates it across the epoch and sends it to your logger and progress bar — no manual metric accumulation.

3. Code in Action

A complete LightningModule

import torch
import torch.nn as nn
import torch.nn.functional as F
import lightning as L

class LitClassifier(L.LightningModule):
    def __init__(self, in_dim=784, hidden=128, num_classes=10, lr=1e-3):
        super().__init__()
        self.save_hyperparameters()               # stores args for checkpointing/logging
        self.net = nn.Sequential(                 # same nn.Module you'd build in Module 03
            nn.Flatten(),
            nn.Linear(in_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, num_classes),
        )
        self.lr = lr

    def forward(self, x):
        return self.net(x)                        # standard forward pass

    def training_step(self, batch, batch_idx):
        x, y = batch                              # batch is ALREADY on the device
        logits = self(x)
        loss = F.cross_entropy(logits, y)         # compute the loss...
        self.log("train_loss", loss)              # ...and log it
        return loss                               # RETURN it — Lightning does backward/step

    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})   # eval()/no_grad() are automatic

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)

The mapping from raw PyTorch (side by side)

# RAW loop body                     ->   LightningModule method
# optimizer.zero_grad()             ->   (automatic)
# logits = model(x.to(device))      ->   training_step: logits = self(x)
# loss = criterion(logits, y)       ->   training_step: loss = F.cross_entropy(...)
# loss.backward()                   ->   (automatic — just `return loss`)
# optimizer.step()                  ->   (automatic)
# optimizer = Adam(model.params())  ->   configure_optimizers
# model.eval()/no_grad() for val    ->   (automatic around validation_step)

Using it is still normal PyTorch

model = LitClassifier()
# It's an nn.Module, so this works exactly as before:
dummy = torch.randn(4, 1, 28, 28)
print(model(dummy).shape)          # torch.Size([4, 10]) -> plain forward pass

4. Common Pitfalls

Pitfall 1 — Forgetting to return loss from training_step. Lightning runs the backward pass on whatever you return. If you compute the loss but don't return it (or return None), no gradients flow and the model silently never learns. Always return the loss (or a dict containing a "loss" key).

Pitfall 2 — Manually calling .to(device), zero_grad, or backward. Inside a LightningModule these are Lightning's job. Doing them yourself double-moves data or double-steps the optimizer, causing errors or corrupted training. Write the step as if the batch is already on-device (it is) and let Lightning handle the mechanics.

Pitfall 3 — Putting metric math in the wrong place or not logging. Accumulating metrics in plain Python variables across steps (the Module 05 habit) fights Lightning's system. Use self.log(...)/self.log_dict(...) so Lightning handles epoch-level aggregation and reduction across devices correctly.

5. Further Reading & Watch List

⬅️ Prev: 01 · Why Lightning? · ➡️ Next: 03 · The Trainer