8  Simplifying the Training Loop

8.1 Introduction

Over the last several chapters, you learned to build a training loop from scratch. This was intentional: understanding what happens in each iteration of the loop (reading a mini-batch, clearing the gradients, predicting, calculating the loss, backpropagating, updating the parameters, and validating at the end of each epoch) is essential for understanding how deep learning models are trained and for troubleshooting when something goes wrong. However, as you saw in Chapter 7, the training loop becomes long and complex as modifications are added, such as learning rate schedulers, gradient accumulation, saving the best model, and logging metrics. Every one of these additions is an opportunity to introduce an error, such as forgetting to reset a metric, calling the scheduler in the wrong place, or allowing the validation data to influence the gradients.

PyTorch Lightning is a library that organizes PyTorch code and handles the training loop for you. You still define the model architecture by subclassing nn.Module, and you still define the loss, metrics, and optimizer. However, Lightning’s Trainer class executes the loop, moves data and models to the correct device, calls zero_grad(), backward(), and step() at the correct times, switches between training and evaluation modes, disables gradient calculation during validation, and logs results. Many of the modifications that required editing the loop in Chapter 7, such as mixed precision training, gradient accumulation, multi-GPU training, and saving the best model, become arguments or callbacks.

In this chapter, I replicate the training process from Chapter 6, in which a fully connected neural network was trained to classify the EuroSAT data using band means, but using PyTorch Lightning as opposed to a manual training loop. I then demonstrate a variety of training loop features that Lightning simplifies. For these features, I only show how they are implemented as opposed to training full models with each of them. From this point forward in the text, we will generally use Lightning to train models.

The R torch ecosystem has a similar package, luz, which is used in our Geospatial Supervised Learning using R text. If you are familiar with luz, you will find that Lightning serves the same purpose.

Lightning is installed with pip install lightning and imported as import lightning as L. Older code may use the pytorch_lightning package name, which refers to the same library. Lightning updates frequently, and some features demonstrated in this chapter, such as weight averaging, were added recently, so you may need to consult the documentation if you are using a different version than the one listed in environment.yml.

8.2 Mapping the Manual Loop to Lightning

The table below shows where each component of our manual training loop from Chapter 6 is defined when using Lightning. The key idea is that you define what happens in a single training or validation step and Lightning handles when it happens.

Manual training loop	PyTorch Lightning
Define architecture by subclassing nn.Module	Same; the model is stored as an attribute of a LightningModule
Instantiate loss, metrics, and optimizer	Loss and metrics in __init__(); optimizer in configure_optimizers()
Loop over epochs and training mini-batches	Handled by Trainer.fit()
Move data and model to the device	Handled by the Trainer
model.train(), zero_grad(), backward(), step()	Handled by the Trainer; you return the loss from training_step()
Predict and calculate loss and metrics for a training mini-batch	training_step()
model.eval() and torch.no_grad() for validation	Handled by the Trainer
Predict and calculate loss and metrics for a validation mini-batch	validation_step()
Accumulate metrics with compute() and reset()	Handled by self.log()
Save losses and metrics to lists and a CSV file	Loggers (e.g., CSVLogger)
Save the model if the validation F1-score improves	ModelCheckpoint callback

8.3 Preparation

I begin by importing the required packages. In addition to the packages used in Chapter 6, I import Lightning, its callbacks and loggers, and its Tuner class.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torchmetrics as tm

import lightning as L
from lightning.pytorch.callbacks import (ModelCheckpoint, EarlyStopping,
                                         LearningRateMonitor)
from lightning.pytorch.loggers import CSVLogger
from lightning.pytorch.tuner import Tuner

The data preparation steps are the same as those used in Chapter 6: I read the tables of band means, calculate the means and standard deviations of the training data for normalization, define the DataSet subclass, and instantiate the DataSets and DataLoaders. Note that I do not need to define a device or move anything to the GPU; the Trainer will handle this.

trainAgg = pd.read_csv("data/files/EuroSat_Tabulated/train_aggregated.csv")
valAgg = pd.read_csv("data/files/EuroSat_Tabulated/val_aggregated.csv")
testAgg = pd.read_csv("data/files/EuroSat_Tabulated/test_aggregated.csv")

trainAggMns = np.array(trainAgg.iloc[:, 3:].mean(axis=0)).flatten()
trainAggSDs = np.array(trainAgg.iloc[:, 3:].std(axis=0)).flatten()

class EuroSat(Dataset):

    def __init__(self, df, bndMns, bndSDs):
        super().__init__()
        self.df = df
        self.bndMns = bndMns
        self.bndSDs = bndSDs

    def __getitem__(self, idx):
        bands = np.array(self.df.iloc[idx, 3:], dtype="float32")
        bands = (bands - self.bndMns) / self.bndSDs
        label = self.df.iloc[idx, 2]
        return torch.from_numpy(bands.astype("float32")), torch.tensor(label).long()

    def __len__(self):
        return len(self.df)

trainDS = EuroSat(trainAgg, trainAggMns, trainAggSDs)
valDS = EuroSat(valAgg, trainAggMns, trainAggSDs)
testDS = EuroSat(testAgg, trainAggMns, trainAggSDs)

trainDL = DataLoader(trainDS, batch_size=256, shuffle=True, drop_last=True)
valDL = DataLoader(valDS, batch_size=256, shuffle=False)
testDL = DataLoader(testDS, batch_size=256, shuffle=False)

The architecture is also unchanged from Chapter 6. This highlights an important point: Lightning does not change how architectures are defined. Any nn.Module subclass, including those provided by torchvision or Segmentation Models, can be trained with Lightning.

class myFCN(nn.Module):
  def __init__(self, inSize, hiddenSizes, outSize):
    super().__init__()
    self.theNetwork = nn.Sequential(
        nn.Linear(inSize, hiddenSizes[0]),
        nn.BatchNorm1d(hiddenSizes[0]),
        nn.ReLU(inplace=True),
        nn.Linear(hiddenSizes[0], hiddenSizes[1]),
        nn.BatchNorm1d(hiddenSizes[1]),
        nn.ReLU(inplace=True),
        nn.Linear(hiddenSizes[1], outSize)
    )

  def forward(self, x):
    return self.theNetwork(x)

8.4 Define a LightningModule

The training process is defined by subclassing L.LightningModule. Here are the key components of the subclass below:

1. In the __init__() constructor method, I instantiate the architecture and store it as the model attribute. Storing the architecture as an attribute, as opposed to defining the layers directly in the LightningModule, keeps the architecture separate from the training process. This makes it easy to reuse the architecture outside of Lightning or to load parameters that were saved from a model trained with a manual loop, as we will see below.
2. I instantiate the loss and store the learning rate as an attribute named lr, which is required by the learning rate finder discussed later in the chapter. self.save_hyperparameters() saves the arguments provided to the constructor so that they are written to the logs and stored in checkpoint files.
3. The assessment metrics are defined using a torchmetrics MetricCollection, which groups several metrics together. Since torchmetrics metrics accumulate results across mini-batches, separate copies must be used for the training, validation, and test data. The clone() method creates these copies, and the prefix argument adds “train_”, “val_”, or “test_” to the metric names in the logs.
4. The forward() method passes the data through the architecture.
5. The training_step() method defines what happens for a single training mini-batch: the batch is unpacked into the predictor variables and targets, the data are predicted, the loss is calculated, the metrics are updated, and the results are logged. The method must return the loss. Lightning then performs the backpropagation and optimization step. There is no need to clear the gradients or move the data to the device.
6. The validation_step() and test_step() methods are similar, but nothing is returned since no parameter updates are performed. Lightning automatically places the model in evaluation mode and disables gradient calculation.
7. The self.log() and self.log_dict() methods log results. Setting on_step=False and on_epoch=True causes the values to be aggregated over the entire epoch. When a torchmetrics metric object is logged, Lightning calls its compute() and reset() methods at the correct times, which removes a common source of error.
8. The configure_optimizers() method returns the optimizer. Note that the optimizer is defined relative to self.parameters().

class LitFCN(L.LightningModule):
    def __init__(self, inSize=10, hiddenSizes=(256, 256), nCls=10, lr=1e-3):
        super().__init__()
        self.save_hyperparameters()
        self.model = myFCN(inSize, list(hiddenSizes), nCls)
        self.criterion = nn.CrossEntropyLoss()
        self.lr = lr

        metrics = tm.MetricCollection({
            "acc": tm.Accuracy(task="multiclass", num_classes=nCls),
            "f1": tm.F1Score(task="multiclass", num_classes=nCls, average="macro"),
            "kappa": tm.CohenKappa(task="multiclass", num_classes=nCls),
        })
        self.train_metrics = metrics.clone(prefix="train_")
        self.val_metrics = metrics.clone(prefix="val_")
        self.test_metrics = metrics.clone(prefix="test_")

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        self.train_metrics.update(logits, y)
        self.log("train_loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        self.log_dict(self.train_metrics, on_step=False, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        self.val_metrics.update(logits, y)
        self.log("val_loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        self.log_dict(self.val_metrics, on_step=False, on_epoch=True, prog_bar=True)

    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        self.test_metrics.update(logits, y)
        self.log("test_loss", loss, on_step=False, on_epoch=True)
        self.log_dict(self.test_metrics, on_step=False, on_epoch=True)

    def predict_step(self, batch, batch_idx):
        x, _ = batch
        return torch.argmax(self(x), dim=1)

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

litModel = LitFCN(inSize=10, hiddenSizes=(256, 256), nCls=10, lr=1e-3)

8.5 Train the Model

The training process is executed using a Trainer. Here, I specify:

1. max_epochs: the number of training epochs (100, as in Chapter 6).
2. accelerator and devices: the hardware to use. "auto" selects a GPU if one is available.
3. logger: where to log results. The CSVLogger writes the logged values to a metrics.csv file inside of a versioned folder (e.g., data/output/eurosat_fcnn/version_0), along with a hparams.yaml file containing the saved hyperparameters.
4. callbacks: a list of callbacks, which are objects that perform actions at specific points in the training process. The ModelCheckpoint callback saves the model when the monitored metric improves, replicating the process used in Chapter 11 of saving the model only when the validation F1-score improves.

If no logger is specified, Lightning writes its logs and checkpoints to a lightning_logs folder in the working directory. The location can be changed using the default_root_dir argument of the Trainer. Throughout this text, all logs, checkpoints, and other files generated by the code are written to the data/output folder, which is empty when you download the book files, so that they do not overwrite the trained models and logs provided in the data/models folder.

The model is then trained by calling fit() and providing the LightningModule and the training and validation DataLoaders.

checkpoint = ModelCheckpoint(monitor="val_f1", mode="max", save_top_k=1,
                             dirpath="data/output/nnEuroSatL/",
                             filename="best-{epoch:02d}-{val_f1:.4f}")

trainer = L.Trainer(
    max_epochs=100,
    accelerator="auto",
    devices=1,
    logger=CSVLogger("data/output/", name="eurosat_fcnn"),
    callbacks=[checkpoint],
)

trainer.fit(litModel, train_dataloaders=trainDL, val_dataloaders=valDL)

That is it! Compare this to the training loop in Chapter 6. The loop itself, the device handling, the gradient clearing, the metric accumulation and resetting, the lists used to store the results, and the code used to save the best model have all been replaced by the LightningModule’s step methods and the Trainer. The path to the best checkpoint file can be obtained from the ModelCheckpoint callback.

checkpoint.best_model_path

8.6 Explore Training Logs

The CSVLogger writes one row to metrics.csv each time values are logged. Since the training and validation metrics are logged at different times within each epoch, they are written to different rows, with missing values in the columns that were not logged in that row. To create a table with one row per epoch, I group the rows by epoch and take the first non-missing value in each column. The result can then be plotted as we did in Chapter 6.

logDF = pd.read_csv(f"{trainer.logger.log_dir}/metrics.csv")
logDF = logDF.groupby("epoch").first().reset_index()
logDF.head()

fig, axs = plt.subplots(1, 2, figsize=(11, 4))
axs[0].plot(logDF["epoch"], logDF["train_loss"], label="Training")
axs[0].plot(logDF["epoch"], logDF["val_loss"], label="Validation")
axs[0].set_title("Loss")
axs[1].plot(logDF["epoch"], logDF["train_f1"], label="Training")
axs[1].plot(logDF["epoch"], logDF["val_f1"], label="Validation")
axs[1].set_title("Macro-Averaged F1-Score")
for ax in axs:
    ax.set_xlabel("Epoch")
    ax.legend()
plt.show()

8.7 Assess the Model

To assess the model using the withheld test data, I call the test() method of the Trainer. Setting ckpt_path="best" loads the parameters from the best checkpoint as opposed to using the parameters from the last epoch, which replaces the process of re-instantiating the model and loading the saved parameters. The results are printed and returned as a list containing a dictionary of the logged test metrics.

testResults = trainer.test(litModel, dataloaders=testDL, ckpt_path="best")
testResults

Predictions can be generated using the predict() method, which calls the predict_step() method defined in the LightningModule for each mini-batch and returns a list of the results.

preds = trainer.predict(litModel, dataloaders=testDL, ckpt_path="best")
preds = torch.cat(preds)

8.8 Saving and Loading Models

Lightning checkpoint files (.ckpt) store the model parameters along with the optimizer state, the saved hyperparameters, and other information required to resume training. A LightningModule can be re-instantiated from a checkpoint using the load_from_checkpoint() class method. Since the hyperparameters were saved using self.save_hyperparameters(), they do not need to be provided again.

bestModel = LitFCN.load_from_checkpoint(checkpoint.best_model_path)
bestModel.eval()

Training can also be resumed from a checkpoint, for example, if training was interrupted, by providing the checkpoint path to fit().

trainer = L.Trainer(max_epochs=150, accelerator="auto", devices=1, default_root_dir="data/output/")
trainer.fit(litModel, train_dataloaders=trainDL, val_dataloaders=valDL,
            ckpt_path=checkpoint.best_model_path)

Since the architecture is stored as the model attribute, the parameters can be extracted and saved as a standard PyTorch state dictionary, which can be loaded into the architecture without Lightning. Similarly, parameters saved from a model trained using a manual training loop, such as the model provided with the text from Chapter 6, can be loaded into the model attribute of the LightningModule. This is how we will use the trained models provided with this text in later chapters.

# Save the architecture's parameters only
torch.save(bestModel.model.state_dict(), "data/output/eurosat_fcnn_lightning.pt")

# Load parameters saved from a manual training loop into the LightningModule
litModel = LitFCN(inSize=10, hiddenSizes=(256, 256), nCls=10)
litModel.model.load_state_dict(torch.load("data/models/nnEurosat.pt"))

8.9 Using a LightningDataModule

Lightning also provides the LightningDataModule class, which collects all of the data preparation steps, including reading the data, calculating the normalization statistics, instantiating the DataSets, and defining the DataLoaders, into a single object. Using a DataModule is optional; as demonstrated above, DataLoaders can be passed directly to fit(). However, a DataModule is useful for organizing and reusing data preparation code, and it is required by the batch size finder discussed below since the finder must be able to rebuild the DataLoaders with different mini-batch sizes. Note that the mini-batch size is stored as the batch_size attribute.

class EuroSatDataModule(L.LightningDataModule):
    def __init__(self, folder="data/files/EuroSat_Tabulated/", batch_size=256, num_workers=0):
        super().__init__()
        self.folder = folder
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage=None):
        trainAgg = pd.read_csv(self.folder + "train_aggregated.csv")
        valAgg = pd.read_csv(self.folder + "val_aggregated.csv")
        testAgg = pd.read_csv(self.folder + "test_aggregated.csv")
        mns = np.array(trainAgg.iloc[:, 3:].mean(axis=0)).flatten()
        sds = np.array(trainAgg.iloc[:, 3:].std(axis=0)).flatten()
        self.trainDS = EuroSat(trainAgg, mns, sds)
        self.valDS = EuroSat(valAgg, mns, sds)
        self.testDS = EuroSat(testAgg, mns, sds)

    def train_dataloader(self):
        return DataLoader(self.trainDS, batch_size=self.batch_size, shuffle=True,
                          drop_last=True, num_workers=self.num_workers)

    def val_dataloader(self):
        return DataLoader(self.valDS, batch_size=self.batch_size, num_workers=self.num_workers)

    def test_dataloader(self):
        return DataLoader(self.testDS, batch_size=self.batch_size, num_workers=self.num_workers)

dm = EuroSatDataModule(batch_size=256)
trainer = L.Trainer(max_epochs=100, accelerator="auto", devices=1, default_root_dir="data/output/")
trainer.fit(litModel, datamodule=dm)
trainer.test(litModel, datamodule=dm, ckpt_path="best")

8.10 Training Loop Features

In the remainder of this chapter, I demonstrate how a variety of training loop features are implemented using Lightning. For each feature, I only show the relevant code; any of them can be combined with the training process demonstrated above by adding the argument or callback to the Trainer.

8.10.1 Model Summaries

Lightning prints a summary of the LightningModule at the beginning of training, which lists each submodule, its type, and its number of parameters, along with the total number of trainable and non-trainable parameters and the estimated model size. The depth of the summary can be controlled using the ModelSummary callback, and a summary can be generated without training using the ModelSummary utility. Unlike torchinfo, which we use in later chapters, Lightning’s summary does not require an input size unless you want the output shapes of each layer. To include them, define an example_input_array attribute in the LightningModule.

from lightning.pytorch.utilities.model_summary import ModelSummary

litModel.example_input_array = torch.zeros(256, 10)  # enables input/output shapes in the summary
print(ModelSummary(litModel, max_depth=3))

from lightning.pytorch.callbacks import ModelSummary as ModelSummaryCallback
trainer = L.Trainer(max_epochs=100, callbacks=[ModelSummaryCallback(max_depth=3)])

8.10.2 Model Checkpoints

The ModelCheckpoint callback was used above to save the model with the highest validation F1-score. Other useful arguments include save_top_k, which saves the best k models (or all models if set to -1); save_last, which also saves the model from the last epoch so that training can be resumed; every_n_epochs, which controls how often checkpoints are considered; and dirpath, which defines the folder in which checkpoints are saved. Multiple ModelCheckpoint callbacks can be used at the same time, for example, to save both the model with the lowest validation loss and the model with the highest validation F1-score.

ckptF1 = ModelCheckpoint(monitor="val_f1", mode="max", save_top_k=3, save_last=True,
                         dirpath="data/output/checkpoints/", filename="f1-{epoch:02d}-{val_f1:.4f}")
ckptLoss = ModelCheckpoint(monitor="val_loss", mode="min", save_top_k=1,
                           dirpath="data/output/checkpoints/", filename="loss-{epoch:02d}-{val_loss:.4f}")

trainer = L.Trainer(max_epochs=100, callbacks=[ckptF1, ckptLoss])

8.10.3 Early Stopping

Early stopping ends the training process when a monitored metric stops improving, which saves time and can reduce overfitting. The EarlyStopping callback monitors a logged metric. The patience argument specifies how many validation checks without improvement are allowed before training stops, and min_delta specifies the minimum change that counts as an improvement. Early stopping is generally combined with a ModelCheckpoint callback so that the best model, as opposed to the model from the last epoch, is retained.

earlyStop = EarlyStopping(monitor="val_loss", mode="min", patience=10, min_delta=0.001)

trainer = L.Trainer(max_epochs=500, callbacks=[earlyStop, checkpoint])

Be careful when combining early stopping with learning rate schedules, such as the one cycle policy, that change the learning rate according to a pre-defined number of epochs. If training stops early, the schedule will not be completed.

8.10.4 Loggers and Experiment Managers

Loggers determine where the logged values are written. In addition to CSVLogger, Lightning supports several experiment managers, which are tools that track and compare training runs, including their hyperparameters, metrics, and saved models. Options include TensorBoard (TensorBoardLogger), MLflow (MLFlowLogger), Weights & Biases (WandbLogger), Comet (CometLogger), and Neptune (NeptuneLogger). TensorBoard and MLflow can be run locally, while Weights & Biases, Comet, and Neptune are hosted services that require an account. The corresponding package must be installed to use each logger.

Multiple loggers can be used at once by providing a list. When comparing many experiments, such as different learning rates or architectures, an experiment manager makes it much easier to keep track of the results than saving and combining many CSV files.

from lightning.pytorch.loggers import TensorBoardLogger, MLFlowLogger

loggers = [
    CSVLogger("data/output/", name="eurosat_fcnn"),
    TensorBoardLogger("data/output/", name="eurosat_fcnn"),
    MLFlowLogger(experiment_name="eurosat_fcnn", tracking_uri="file:./mlruns"),
]
trainer = L.Trainer(max_epochs=100, logger=loggers)

TensorBoard can then be launched from the command line using tensorboard --logdir data/output, and MLflow using mlflow ui. Both provide interactive web-based dashboards for comparing runs.

8.10.5 Mixed and 16-Bit Precision

By default, PyTorch uses 32-bit floating point numbers for the model parameters and calculations. Mixed precision training performs many operations, such as convolutions and matrix multiplications, using 16-bit floating point numbers while keeping operations that require higher precision, such as the parameter updates and some reductions, in 32-bit. This reduces memory usage, which allows for larger mini-batches or larger models, and can greatly speed up training on GPUs with hardware support for 16-bit computation (e.g., the Tensor Cores in modern NVIDIA GPUs).

There are two 16-bit formats. The standard half-precision format (float16) has a limited range, so very small gradients can become zero (underflow). To prevent this, the loss is multiplied by a scale factor before backpropagation and the gradients are divided by the same factor before the update, a process called gradient scaling. The bfloat16 format has the same range as float32 but lower precision, so gradient scaling is not required. bfloat16 is supported by NVIDIA GPUs with the Ampere architecture or newer (e.g., RTX 30-series GPUs and newer).

Implementing mixed precision training in a manual loop requires wrapping the forward pass in torch.autocast() and using a GradScaler for the backward pass and optimizer step. With Lightning, it only requires setting the precision argument.

trainer = L.Trainer(max_epochs=100, precision="16-mixed")    # float16 with gradient scaling
trainer = L.Trainer(max_epochs=100, precision="bf16-mixed")  # bfloat16 (Ampere or newer GPUs)

Mixed precision generally has little impact on model performance, but it is a good idea to compare the results to a 32-bit model when you first adopt it for a new problem. Some custom loss functions may be sensitive to reduced precision, so it may be necessary to calculate them in 32-bit by converting the inputs using .float().

8.10.6 Multi-GPU Training

In Chapter 7, we used nn.DataParallel() to train using multiple GPUs. The recommended approach in PyTorch is distributed data parallel (DDP), in which a separate process is run for each GPU, each process holds a copy of the model and processes a different portion of each mini-batch, and the gradients are averaged across processes before each parameter update. DDP is faster than nn.DataParallel() but more complicated to set up manually since it requires launching and coordinating multiple processes and using a distributed sampler so that each process receives different data. With Lightning, it only requires setting the devices and strategy arguments.

trainer = L.Trainer(max_epochs=100, accelerator="gpu", devices=2, strategy="ddp")

A few notes about multi-GPU training:

1. The mini-batch size defined in the DataLoader is the mini-batch size per GPU, so the effective batch size is the batch size multiplied by the number of GPUs. You may need to adjust the learning rate accordingly.
2. When logging values that are not torchmetrics metrics, set sync_dist=True in self.log() so that the values are averaged across GPUs. torchmetrics metrics are synchronized automatically.
3. DDP launches the script multiple times, so it works best when training is run from a Python script as opposed to an interactive notebook. For notebooks, Lightning provides the "ddp_notebook" strategy. On Windows, the NCCL backend used for communication between GPUs on Linux is not available, so Lightning uses the slower Gloo backend.

8.10.7 Gradient Accumulation and Clipping

Gradient accumulation, which required editing the training loop in Chapter 7, is implemented with the accumulate_grad_batches argument. Gradient clipping, which limits the magnitude of the gradients to prevent very large parameter updates and unstable training, is implemented with the gradient_clip_val argument.

trainer = L.Trainer(max_epochs=100, accumulate_grad_batches=4, gradient_clip_val=1.0)

8.10.8 Learning Rate Finder

Lightning provides its own implementation of the learning rate range test discussed in Chapter 7. The Tuner class wraps a Trainer, and its lr_find() method trains the model for a small number of iterations while increasing the learning rate exponentially, recording the loss at each learning rate. The results can be plotted, and a suggested learning rate, at the point with the steepest decrease in the loss, is provided. Setting update_attr=True sets the LightningModule’s lr attribute to the suggested value. This is why the learning rate was stored as self.lr in the LightningModule and used in configure_optimizers().

trainer = L.Trainer(max_epochs=100, accelerator="auto", devices=1)
tuner = Tuner(trainer)
lrFinder = tuner.lr_find(litModel, train_dataloaders=trainDL, val_dataloaders=valDL,
                         min_lr=1e-6, max_lr=1, num_training=100, update_attr=True)
fig = lrFinder.plot(suggest=True)
plt.show()
print(lrFinder.suggestion())

As with any learning rate finder, the suggestion should be treated as a starting point. I recommend inspecting the plot as opposed to relying only on the suggested value.

8.10.9 Mini-Batch Size Finder

The batch size finder searches for the largest mini-batch size that fits into the available memory. Starting from the current mini-batch size, it repeatedly runs a few training steps and doubles the mini-batch size (mode="power") until an out-of-memory error occurs, then selects the largest mini-batch size that worked. mode="binsearch" refines the result using a binary search. The mini-batch size must be stored as a batch_size attribute of the DataModule (or LightningModule), and the DataLoaders must be created using that attribute, since the finder needs to rebuild the DataLoaders with each candidate mini-batch size.

dm = EuroSatDataModule(batch_size=32)
trainer = L.Trainer(max_epochs=100, accelerator="auto", devices=1)
tuner = Tuner(trainer)
tuner.scale_batch_size(litModel, datamodule=dm, mode="power")
print(dm.batch_size)

Note that the largest mini-batch size that fits in memory is not necessarily the best mini-batch size for model performance. Very large mini-batch sizes reduce the number of parameter updates per epoch and can result in poorer generalization unless the learning rate is adjusted. The batch size finder is most useful when working with large images or large models where memory, as opposed to model performance, is the limiting factor.

8.10.10 Learning Rate Schedulers

Learning rate schedulers are returned from configure_optimizers() along with the optimizer. The "interval" key specifies whether the scheduler steps after each mini-batch ("step") or each epoch ("epoch"). In Chapter 7, we had to remember to call scheduler.step() after each mini-batch for the one cycle policy; with Lightning, setting "interval": "step" handles this. The estimated_stepping_batches property of the Trainer provides the total number of optimization steps, which is needed by the one cycle scheduler and accounts for gradient accumulation and multiple GPUs. The LearningRateMonitor callback logs the learning rate so that the schedule can be verified.

class LitFCNOneCycle(LitFCN):
    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
        scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer, max_lr=0.01, total_steps=self.trainer.estimated_stepping_batches)
        return {"optimizer": optimizer,
                "lr_scheduler": {"scheduler": scheduler, "interval": "step"}}

trainer = L.Trainer(max_epochs=100, callbacks=[LearningRateMonitor(logging_interval="step")])

Schedulers that respond to a metric, such as ReduceLROnPlateau, which reduces the learning rate when the validation loss stops improving, require a "monitor" key.

def configure_optimizers(self):
    optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=5)
    return {"optimizer": optimizer,
            "lr_scheduler": {"scheduler": scheduler, "monitor": "val_loss", "interval": "epoch"}}

8.10.11 Exponential Moving Average of Weights

An exponential moving average (EMA) of the model parameters is a technique in which a second copy of the parameters is maintained during training and updated after each optimization step as a weighted average of its previous value and the current parameters:

\[\theta_{EMA} \leftarrow \lambda \, \theta_{EMA} + (1 - \lambda) \, \theta\]

where \(\theta\) represents the current parameters and \(\lambda\), the decay, is generally close to 1 (e.g., 0.999). The averaged parameters are not used to calculate the gradients; they simply track a smoothed version of the parameters over time. Since stochastic gradient descent causes the parameters to fluctuate from one mini-batch to the next, especially with large learning rates or small mini-batches, the averaged parameters often generalize better and produce more stable validation results than the parameters from any single step. The averaged parameters are used for validation, testing, and inference. EMA is widely used for training modern architectures, including many transformer and ConvNeXt models, and it is also a key component of some self-supervised learning methods.

A related method, stochastic weight averaging (SWA), computes a simple average of the parameters from the later epochs of training, often combined with a constant or cyclical learning rate. Lightning has long provided a StochasticWeightAveraging callback. Newer versions of Lightning provide a more general WeightAveraging callback, which uses PyTorch’s AveragedModel class and can be configured to compute an EMA. In the example below, I subclass WeightAveraging to define an EMA with a decay of 0.999 that begins updating after the first 100 optimization steps, since the parameters change rapidly early in training.

from lightning.pytorch.callbacks import WeightAveraging
from torch.optim.swa_utils import get_ema_multi_avg_fn

class EMAWeightAveraging(WeightAveraging):
    def __init__(self, decay=0.999, startStep=100):
        super().__init__(multi_avg_fn=get_ema_multi_avg_fn(decay))
        self.startStep = startStep

    def should_update(self, step_idx=None, epoch_idx=None):
        return (step_idx is not None) and (step_idx >= self.startStep)

trainer = L.Trainer(max_epochs=100, callbacks=[EMAWeightAveraging(decay=0.999), checkpoint])

from lightning.pytorch.callbacks import StochasticWeightAveraging
trainer = L.Trainer(max_epochs=100, callbacks=[StochasticWeightAveraging(swa_lrs=1e-3, swa_epoch_start=0.8)])

If you are using an older version of Lightning without the WeightAveraging callback, or a manual training loop, an EMA can be implemented using PyTorch’s AveragedModel class directly. The averaged model is updated after each optimizer step and is used for validation. Note that if the model contains batch normalization layers, use_buffers=True should be used so that the running mean and variance used by batch normalization are also averaged.

from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn

emaModel = AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(0.999), use_buffers=True)

# Inside the training loop, after optimizer.step():
emaModel.update_parameters(model)

# Use emaModel (as opposed to model) for validation and inference

8.10.12 Reproducibility

The settings discussed in Chapter 7 for making training reproducible are simplified by Lightning’s seed_everything() function, which seeds the Python, NumPy, and PyTorch generators, and the deterministic argument of the Trainer, which enables deterministic algorithms. Setting workers=True also seeds each DataLoader worker. The same limitations discussed in Chapter 7 still apply.

L.seed_everything(42, workers=True)
trainer = L.Trainer(max_epochs=100, deterministic=True)

8.10.13 Debugging Utilities

Lightning provides several Trainer arguments that are useful for quickly checking that a training process works before committing to a long training run. fast_dev_run=True runs a single training and validation mini-batch, which will quickly reveal errors in the DataSet, architecture, or step methods. overfit_batches trains on a small number of mini-batches repeatedly; if the model cannot overfit to a single mini-batch, there is likely a problem with the architecture, loss, or data. limit_train_batches and limit_val_batches use only a fraction of the data each epoch. The profiler argument reports the time spent in each part of the training process, which helps identify bottlenecks, such as slow data loading.

trainer = L.Trainer(fast_dev_run=True)                       # one batch, catches errors quickly
trainer = L.Trainer(max_epochs=50, overfit_batches=1)        # should drive training loss toward 0
trainer = L.Trainer(max_epochs=5, limit_train_batches=0.1, limit_val_batches=0.1)
trainer = L.Trainer(max_epochs=1, profiler="simple")

8.10.14 Combining Features

As a final example, the code below combines many of the features discussed in this chapter into a single training configuration. Note how each feature is added as a single argument or callback without changing the LightningModule.

L.seed_everything(42, workers=True)

trainer = L.Trainer(
    max_epochs=200,
    accelerator="gpu",
    devices=1,
    precision="bf16-mixed",
    accumulate_grad_batches=2,
    gradient_clip_val=1.0,
    logger=[CSVLogger("data/output/", name="eurosat_fcnn"),
            TensorBoardLogger("data/output/", name="eurosat_fcnn")],
    callbacks=[
        ModelCheckpoint(monitor="val_f1", mode="max", save_top_k=1, save_last=True),
        EarlyStopping(monitor="val_loss", mode="min", patience=15),
        LearningRateMonitor(logging_interval="epoch"),
        EMAWeightAveraging(decay=0.999),
    ],
)
trainer.fit(litModel, datamodule=dm)
trainer.test(litModel, datamodule=dm, ckpt_path="best")

8.11 Concluding Remarks

PyTorch Lightning removes much of the repetitive and error-prone code associated with training loops while still allowing you to define the architecture, loss, metrics, and optimizer using standard PyTorch. It also makes features such as mixed precision, multi-GPU training, checkpointing, early stopping, experiment tracking, and weight averaging simple to implement. However, Lightning is not a substitute for understanding the training loop. When something goes wrong, knowing what Lightning is doing on your behalf, as explored in the prior chapters, is essential for diagnosing the problem.

This chapter concludes the first part of the text. In the next part, we will explore convolutional neural networks, which allow the spatial context of image data to be incorporated into the learning process. From this point forward, we will generally use Lightning to train models.

8.12 Questions

1. Explain the roles of the LightningModule and the Trainer. Which components of a manual training loop are defined in the LightningModule, and which are handled by the Trainer?
2. What must the training_step() method return, and why? Which operations from the manual training loop no longer need to be written explicitly?
3. Why must separate copies of each torchmetrics metric be used for the training, validation, and test data? How does self.log() simplify the accumulation and resetting of metrics?
4. Explain the purposes of the ModelCheckpoint and EarlyStopping callbacks. Why are they generally used together?
5. Explain the difference between float16 and bfloat16 mixed precision training. Why does float16 training require gradient scaling?
6. Explain how distributed data parallel (DDP) training works. How does the number of GPUs affect the effective mini-batch size?
7. Explain how an exponential moving average (EMA) of the model parameters is calculated and why the averaged parameters may generalize better than the parameters from the final optimization step.
8. What attributes must be defined for Lightning’s learning rate finder and mini-batch size finder to work? Why might the largest mini-batch size that fits in memory not be the best batch size for model performance?

8.13 Exercises

Exercise 8.1: Training with PyTorch Lightning

Introduction

The goal of this exercise is to re-implement the training process from Exercise 6.1 using PyTorch Lightning and to experiment with several of the training loop features that Lightning simplifies. You will work with the same Covertype dataset (covertype.csv in the data/files/tables folder), forest type classes, predictor variables, and data partitions used in Exercise 6.1.

Deliverables

• A Quarto (.qmd) file containing all required code and answers to the questions and discussion prompts, and the HTML webpage rendered from it

Tasks

Task 1: Define a LightningDataModule that reads the Covertype data, partitions it into training (60%), validation (20%), and testing (20%) sets, normalizes the predictor variables to z-scores using statistics from the training data only, and returns DataLoaders with a batch_size attribute.

Task 2: Define a LightningModule that uses the fully connected architecture from Exercise 6.1, cross entropy loss, and the AdamW optimizer and that logs the loss, overall accuracy, and class-aggregated macro-averaged F1-score for the training, validation, and testing data using torchmetrics.

Task 3: Use Lightning’s learning rate finder to select a learning rate. Include the learning rate finder plot and explain your choice.

Task 4: Train the model for up to 100 epochs using a ModelCheckpoint callback that saves the model with the largest validation F1-score, an EarlyStopping callback that monitors the validation loss with a patience of 10 epochs, and a CSVLogger. Plot the training and validation losses and F1-scores by epoch.

Task 5: Re-train the model using bfloat16 or float16 mixed precision and an exponential moving average of the model weights. Compare the training time, loss curves, and test-set metrics to those from Task 4.

Task 6: Assess the best model from Task 4 using the test data and the test() method. Discuss how the results compare to those you obtained in Exercise 6.1 and the advantages and disadvantages of using Lightning as opposed to a manual training loop.