Training a Neural Network#
Initially, right after the neural network is created, the parameters of all layers are filled with small random values. This step is called random initialization. At this point, the network merely implements a series of random transformations. The next step is to gradually adjust these parameters based on the available data. This process is called training and consists of repeating the following steps as long as necessary.
Data sampling: A batch of data (inputs and targets) are randomly selected from the training set.
Forward pass: The inputs are passed through the network, and the outputs are computed.
Loss computation: The mismatch between the network outputs and the targets is measured.
Backward pass: The gradients of the loss with respect to the network parameters are computed.
Parameters update: The network parameters are updated using the computed gradients.
Eventually, the network learns to make accurate predictions on the training data by minimizing the loss function. In this tutorial, we will explain how to implement the training process using PyTorch.
Preparation#
Building a deep learning pipeline starts with a task and a dataset. Here we classify MNIST digits. The dataset contains 60,000 development images and 10,000 test images of size 28×28. We split the development partition into training and validation subsets. Validation allows us to tune hyperparameters and decide when to stop training. The test set remains untouched until the model is final.
Dataset#
We load the MNIST dataset with a preprocessing pipeline that converts the images into PyTorch tensors and rescales them to have values between 0 and 1.
preprocess = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
full_ds = MNIST('.data', train=True, download=True, transform=preprocess)
test_ds = MNIST('.data', train=False, download=True, transform=preprocess)
train_ds, valid_ds = torch.utils.data.random_split(
full_ds, [0.9, 0.1], generator=torch.Generator().manual_seed(42)
)
Neural network#
For classifying the MNIST digits, we define a simple feedforward neural network that flattens the inputs and passes them through two fully-connected layers. The first layer has 512 neurons and a ReLU activation. The second layer has no activation since we will use the cross-entropy loss, which includes a softmax activation. The number of input in the first layer and the number of outputs in the last layer are specified as arguments to the constructor.
class SimpleNet(nn.Module):
def __init__(self, input_dim, num_classes):
super().__init__()
self.flatten = nn.Flatten()
self.linear1 = nn.Linear(input_dim, 512)
self.linear2 = nn.Linear(512, num_classes)
def forward(self, x):
y = self.flatten(x)
y = self.linear1(y)
y = F.relu(y)
y = self.linear2(y)
return y
Loss function#
To control the output of a neural network, we need to be able to measure how far this output is from what we expected. This is the job of the loss function. It takes the prediction of the network and the expected target (what you wanted the network to output), and computes a distance score, capturing how well the network has done on this specific sample. The loss function is a key component of the training process, as it guides the optimization algorithm to adjust the network’s parameters in the right direction.
PyTorch provides a list of predefined loss functions. Choosing the right loss function for the right problem is extremely important, as a network will take any shortcut it can to minimize the loss. Fortunately, when it comes to common problems such as classification and regression, there are simple guidelines we can follow to choose the correct loss function.
Binary cross-entropy for a two-class classification.
Categorical cross-entropy for a multi-class classification problem.
Mean squared error for a regression problem.
Handwritten digit classification is a multi-class classification problem, so we will use the categorical cross-entropy loss function. This is called nn.CrossEntropyLoss in PyTorch.
loss_fn = nn.CrossEntropyLoss()
Note
You don’t need to include a softmax activation in the network when using nn.CrossEntropyLoss, as this function computes the softmax and the cross-entropy loss together. Conversely, if you are using nn.LogSoftmax as the output activation, then you should use nn.NLLLoss instead.
Optimizer#
The central idea behind deep learning is to adjust the parameters of a neural network using the gradient of the loss function. This is possible because the gradient is basically a vector that tells us in which direction we should move each parameter to reduce the loss. The optimizer is the algorithm responsible for adjusting the network’s parameters based on the computed gradients. The most common optimizer is stochastic gradient descent (SGD), but there are many different optimizers available in PyTorch, such as ADAM and RMSProp, that work better for different kinds of models and data.
To construct an Optimizer, we have to give it an iterable containing the parameters to optimize. In this case, we provide it with all the network parameters, which can be iterated over using the parameters() method. We can also specify optimizer-specific options, such as the learning rate.
model = SimpleNet(28*28, 10)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
Note
The parameters to optimize must be tensors that have their requires_grad attribute set to True. All trainable parameters in a PyTorch model have this attribute set by default.
Training loop#
At this point, we have all the pieces to start training our neural network. But here comes the tricky part: how do we put all these pieces together? PyTorch does not answer this question for us, since it is designed as a low-level library that provides the building blocks for deep learning, but does not impose any specific way to use them. Over the years, several third-party libraries have been developed to provide high-level APIs for training neural networks, such as PyTorch Lightning. These libraries can save time and reduce boilerplate. However, it is instructive to understand how training works under the hood.
Note
A custom training loop is defined in the training.py script. Download this file to your working directory to reproduce the examples presented in these pages.
Implementation#
A good way to structure a training loop is to separate the per-batch logic from the boilerplate code that drives the loop itself. Conceptually, we can break down the implementation into two functions.
train_step(): This function will update the network’s parameters on a single batch of data.Trainer: This class will manage the overall training process. It provides a.fit()method that iterates over the dataset in batches and callstrain_step()to perform the actual training work for each batch.
We will now look at a minimal implementation of these two components.
The code snippet below shows the function that trains the model on a single batch of data. The assumption is that the batch is a tuple containing the inputs and the targets, the model takes the inputs as its only argument, and the loss function takes the model outputs and the targets as its arguments.
def train_step(model: nn.Module,
batch: tuple[torch.Tensor, torch.Tensor],
loss_fn: nn.Module,
optimizer: optim.Optimizer,
device: torch.device):
# Data transfer
inputs, labels = batch
inputs = inputs.to(device)
labels = labels.to(device)
# Clear gradients left by the previous mini-batch
optimizer.zero_grad()
# Forward pass
outputs = model(inputs)
# Backward pass
loss = loss_fn(outputs, labels)
loss.backward()
# Model update
optimizer.step()
return loss.item()
The next code snippet shows the class that manages the training process. Notice that the fit() method delegates the processing of each batch to the function defined above. This separation allows for greater flexibility, as the training loop can be reused by replacing only the train_step() function.
class Trainer:
def fit(self,
model: nn.Module,
loader: DataLoader,
loss_fn: nn.Module,
optimizer: optim.Optimizer,
epochs: int):
# Device transfer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# Data iteration
for epoch in range(epochs):
model.train()
for batch in loader:
train_step(model, batch, loss_fn, optimizer, device)
This minimal loop is intentionally small. A practical loop also averages loss with the correct sample weighting, switches between model.train() and model.eval(), evaluates under torch.inference_mode(), logs useful diagnostics, and saves the checkpoint with the best validation result.
Note
A powerful debugging ritual is to train on just one small batch. A sufficiently expressive model should nearly memorize it. If the loss cannot become very small, inspect shapes, targets, activations, gradient clearing, and the learning rate before launching a full run.
Usage example#
Let’s demonstrate how to use the Trainer class. We start by importing it from the training.py file.
from training import Trainer
Next, we train for up to 10 epochs, using batches of 512 and a learning rate of \(10^{-3}\). Large batches give smoother gradient estimates but perform fewer updates per epoch and use more memory; small batches are noisier and may generalize differently. These are validation choices, not universal constants.
After each epoch, we compare validation loss and keep an in-memory copy of the best parameters. This is a simple form of checkpoint selection.
model = SimpleNet(28*28, 10)
train_loader = DataLoader(train_ds, batch_size=512, shuffle=True)
valid_loader = DataLoader(valid_ds, batch_size=512, shuffle=False)
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
epochs = 10
trainer = Trainer()
history = trainer.fit(model, train_loader, loss_fn, optimizer, epochs, valid_loader)
===== Training on cpu device =====
Epoch 1/10: 100%|██████████| 106/106 [00:12<00:00, 8.44it/s, train_loss=0.5220, valid_loss=0.2721]
Epoch 2/10: 100%|██████████| 106/106 [00:12<00:00, 8.66it/s, train_loss=0.2224, valid_loss=0.2017]
Epoch 3/10: 100%|██████████| 106/106 [00:12<00:00, 8.17it/s, train_loss=0.1648, valid_loss=0.1643]
Epoch 4/10: 100%|██████████| 106/106 [00:12<00:00, 8.18it/s, train_loss=0.1278, valid_loss=0.1374]
Epoch 5/10: 100%|██████████| 106/106 [00:12<00:00, 8.44it/s, train_loss=0.1023, valid_loss=0.1212]
Epoch 6/10: 100%|██████████| 106/106 [00:11<00:00, 9.55it/s, train_loss=0.0835, valid_loss=0.1081]
Epoch 7/10: 100%|██████████| 106/106 [00:11<00:00, 8.94it/s, train_loss=0.0697, valid_loss=0.1007]
Epoch 8/10: 100%|██████████| 106/106 [00:12<00:00, 8.75it/s, train_loss=0.0587, valid_loss=0.0950]
Epoch 9/10: 100%|██████████| 106/106 [00:13<00:00, 7.98it/s, train_loss=0.0502, valid_loss=0.0937]
Epoch 10/10: 100%|██████████| 106/106 [00:14<00:00, 7.28it/s, train_loss=0.0432, valid_loss=0.0862]
The curves tell two related stories. Training loss shows whether optimization is working; validation loss shows whether that progress transfers to unseen samples.
Note
A flat training curve may indicate a bad learning rate, broken gradients, or mismatched targets.
A falling training loss paired with a rising validation loss is the classic signature of overfitting.
Inference#
Now that the model is trained, we can use it to make predictions on new data. Let’s select a few images from the test set and ask the model to predict the digit in each one. We will also plot the images to see what the model is working with.
Summary#
In this tutorial, we learned how to train a neural network using PyTorch. Specifically, we discussed how to choose a loss function, select an optimizer, and implement the training loop. We also learned how to make predictions with the trained model. Next, we will discuss how to evaluate the trained model.