03 · The Trainer 🎛️

4 min read

The LightningModule says what to do; the Trainer does it. This one object replaces every loop you've ever written — and unlocks GPUs, mixed precision, and logging by flipping a flag.

1. The 'Why'

With your model reorganized into a LightningModule (sub-module 02), something is conspicuously missing: the loop that actually runs it. That's deliberate — the loop is exactly the boilerplate Lightning promised to take away, and the Trainer is where it now lives. You hand the Trainer your LightningModule and your DataLoaders, call .fit(...), and it executes the entire training process: iterating epochs and batches, calling your training_step and validation_step, running zero_grad/backward/step, switching modes, moving data to the device, and driving the progress bar. The dozens of lines you wrote by hand in Module 05 collapse into essentially one call.

The deeper win is what the Trainer unlocks for free. Because Lightning owns the loop, it can inject powerful engineering features that would each be painful to hand-code correctly. Want to train on a GPU instead of CPU? Change one argument. Want mixed-precision (16-bit) training to run roughly twice as fast and use less memory? One argument. Multiple GPUs, gradient accumulation, a quick debugging run on a handful of batches, automatic logging to TensorBoard? All arguments to the same Trainer. This is the concrete return on the discipline of restructuring your code — you write your model once, and scaling from a laptop CPU to a multi-GPU server becomes a configuration change rather than a rewrite.

2. Core Concepts

Trainer runs everything. You configure it once (Trainer(...)), then call trainer.fit(model, train_loader, val_loader). It owns the epoch/batch loops and calls your LightningModule hooks at the right moments.

.fit, .validate, .test, .predict. fit trains (with validation); validate/test run those splits on a trained model; predict runs inference. Each maps to the corresponding *_step in your module.

Hardware & precision by flag. accelerator="gpu", devices=2, precision="16-mixed" — Lightning handles device placement, distribution, and autocasting. No manual .to(device) or AMP code.

Fast iteration flags. fast_dev_run=True runs a single batch to smoke-test your code; limit_train_batches=0.1 trains on 10% for a quick sanity check; max_epochs caps training length.

Logging is built in. self.log(...) calls in your module flow to a logger (TensorBoard by default) and the progress bar automatically — the metric tracking you did by hand in Module 05 is now free.

3. Code in Action

The whole training run in a few lines

import lightning as L
from torch.utils.data import DataLoader

model = LitClassifier()                          # from sub-module 02
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
val_loader   = DataLoader(val_ds,   batch_size=64)

trainer = L.Trainer(max_epochs=10)               # configure the loop
trainer.fit(model, train_loader, val_loader)     # ...and run it — that's it

Scaling to hardware and precision by changing arguments

import lightning as L

trainer = L.Trainer(
    max_epochs=10,
    accelerator="auto",     # use GPU/MPS if available, else CPU — device-agnostic
    devices="auto",         # use all available devices
    precision="16-mixed",   # mixed precision: ~2x faster, less memory, one flag
)
trainer.fit(model, train_loader, val_loader)
# No manual .to(device), no autocast/GradScaler code — Lightning handles it.

Fast debugging and final evaluation

import lightning as L

# Smoke-test the whole pipeline on ONE batch before a long run
debug = L.Trainer(fast_dev_run=True)
debug.fit(model, train_loader, val_loader)       # catches bugs in seconds

# After training, evaluate on the test set (calls your test_step / validation_step)
trainer = L.Trainer(max_epochs=10)
trainer.fit(model, train_loader, val_loader)
trainer.test(model, test_loader)                 # reports logged test metrics

4. Common Pitfalls

Pitfall 1 — Re-adding manual device/precision code. Once the Trainer manages hardware, sprinkling .to(device) or hand-rolled AMP back into your module fights the framework and causes device errors. Configure hardware on the Trainer and keep the LightningModule device-agnostic.

Pitfall 2 — Skipping fast_dev_run before a long job. Launching a multi-hour run only to crash at the validation step wastes time. Run Trainer(fast_dev_run=True) first — it executes a single train/val/test batch and surfaces shape and logic bugs in seconds.

Pitfall 3 — Expecting to still call .backward() or loop yourself. Beginners sometimes wrap trainer.fit in their own for epoch loop or call optimizer methods manually. The Trainer is the loop; let it own epochs and optimization. Control training length with max_epochs, not an external loop.

5. Further Reading & Watch List

📖 Deep Learning with PyTorch + Lightning docs to scale training
🌐 Lightning — Trainer
🛠️ Trainer flags reference · Mixed precision training
▶️ Lightning AI — official tutorials & walkthroughs

⬅️ Prev: 02 · The LightningModule · ➡️ Next: 04 · DataModules & Callbacks