For our example, we use the Fashion-MNIST dataset, although the specific dataset is not particularly important for demonstrating these concepts.
Initial Setup
We begin by importing the required libraries and loading our dataset. The model is configured to train for eight epochs, which represents the maximum number of epochs the training process is allowed to run.
However, thanks to early stopping, the training may end before reaching this limit.
Implementing Early Stopping in PyTorch
Since PyTorch does not provide built-in early stopping functionality, we must implement it manually within the training loop.
Tracking the Best Validation Loss
Early stopping is based on monitoring the validation loss. We start by initializing the best validation loss to infinity. This guarantees that the first computed validation loss will always be an improvement.
Defining an Improvement Threshold
Not every improvement is meaningful. To avoid reacting to very small fluctuations, we define a minimum improvement threshold equal to 0.001. If the validation loss fails to improve by at least this value, the epoch is considered below the performance threshold.
Setting Patience
The patience parameter determines how many times the validation loss is allowed to fall below the threshold before training is stopped. In this example, patience is set to two, meaning that training will stop early if the validation loss does not sufficiently improve two times.
Early Stopping Logic in the Training Loop
At the end of each epoch, we compute the validation loss and calculate the difference between the best validation loss and the current validation loss. If the current validation loss is lower, we update the best value. If the improvement does not meet the threshold, a counter is increased. When this counter reaches the patience value, the training loop is interrupted.
This logic allows the model to stop training once it stops learning meaningful patterns.
Dropout for Regularization
Another important technique covered in this demo is dropout, which helps prevent overfitting. In PyTorch, dropout is implemented as a layer, where the parameter p represents the probability that a neuron is zeroed out.
With p set to 0.5, half of the inputs to the dropout layer are randomly set to zero during training. One advantage of PyTorch’s implementation is that the same dropout layer can be reused throughout the model without worrying about input or output sizes.
Regularization and Momentum in the Optimizer
Both L2 regularization and momentum are implemented directly in the optimizer. Regularization is controlled using the weight decay parameter, while momentum is controlled using the momentum parameter.
When using stochastic gradient descent, these options help stabilize training and improve the model’s ability to generalize.
Putting Everything Together
All of these techniques—early stopping, dropout, regularization, and momentum—can be used simultaneously in a single training pipeline.
After training the model, we observe that all eight epochs are completed, even though early stopping is enabled. This happens because the validation loss only fails to improve beyond the threshold in Epoch 6 and Epoch 8. Since the patience value is two, training is allowed to continue.
If the model had been configured to train for ten epochs instead of eight, early stopping would have triggered at Epoch 8