Early Stopping in PyTorch: Preventing Overfitting During Training

Introduction: Why Use Early Stopping?

Welcome back! In the last lesson, you learned how to use dropout to help your neural network generalize better and avoid overfitting. As a quick reminder, overfitting happens when your model learns the training data too well, including its noise, and then struggles to perform well on new, unseen data. Dropout is one way to address this, but it is not the only tool available. In this lesson, I will introduce you to another important technique: early stopping . Early stopping is a simple but powerful method to prevent overfitting during training. Instead of training your model for a fixed number of epochs, you monitor its performance on a validation set and stop training when the model stops improving. This way, you avoid wasting time and resources on training that does not help your model get better — and you also reduce the risk of overfitting. Early stopping is widely used in deep learning and is especially helpful when you are not sure how many epochs your model really needs.

How Early Stopping Works

The main idea behind early stopping is to keep an eye on your model’s performance on a validation set during training. After each epoch, you check the validation loss — a measure of how well your model is doing on data it has not seen before. If the validation loss keeps getting better, you continue training. But if the validation loss stops improving for a certain number of epochs, called the patience, you stop training early. Patience is a key parameter in early stopping. It tells your training loop how many epochs to wait for an improvement before giving up. For example, if patience is set to 5 and the validation loss does not improve for 5 epochs in a row, training will stop. This helps you avoid stopping too soon if there is a small bump in the loss, but also prevents you from training for too long when there is no real progress.

Best Model vs. Last Model

One important detail to keep in mind: the model parameters at the end of training (the "last model") may not correspond to the best performance on the validation set. Often, the best model (the one with the lowest validation loss) occurs several epochs before training actually stops. If you only use the model as it is at the end of training, you might not get the best results. To address this, you should save a copy of the model’s parameters whenever a new best validation loss is achieved. After early stopping triggers, you can reload these saved parameters to ensure you are using the best version of your model.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Example: Adding Early Stopping and Saving the Best Model in PyTorch

Let’s look at how you can add early stopping to your PyTorch training loop and also save the best model. Below is a code example that shows a typical way to implement early stopping and model checkpointing. This code assumes you already have your model, optimizer, loss function, and data prepared, just like in the previous lessons.

Python

best_loss = float('inf')
patience, wait = 5, 0
best_model_state = None

for epoch in range(100):
    model.train()
    optimizer.zero_grad()
    output = model(X_train)
    loss = loss_fn(output, y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        val_loss = loss_fn(model(X_val), y_val)
        print(f"Epoch {epoch}: Val Loss = {val_loss.item():.4f}")
        if val_loss < best_loss:
            best_loss = val_loss
            wait = 0
            best_model_state = model.state_dict()  # Save the best model
        else:
            wait += 1
            if wait >= patience:
                print("Early stopping!")
                break

# After training, load the best model weights
if best_model_state is not None:
    model.load_state_dict(best_model_state)

best_loss = float('inf')
patience, wait = 5, 0
best_model_state = None

for epoch in range(100):
    model.train()
    optimizer.zero_grad()
    output = model(X_train)
    loss = loss_fn(output, y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        val_loss = loss_fn(model(X_val), y_val)
        print(f"Epoch {epoch}: Val Loss = {val_loss.item():.4f}")
        if val_loss < best_loss:
            best_loss = val_loss
            wait = 0
            best_model_state = model.state_dict()  # Save the best model
        else:
            wait += 1
            if wait >= patience:
                print("Early stopping!")
                break

# After training, load the best model weights
if best_model_state is not None:
    model.load_state_dict(best_model_state)

Here’s how this code works. At the start, best_loss is set to infinity, so any real validation loss will be better. The patience variable is set to 5, which means the training will stop if the validation loss does not improve for 5 epochs in a row. The wait variable keeps track of how many epochs have passed since the last improvement. The best_model_state variable is used to store the parameters of the best model seen so far.

During each epoch, the model is trained on the training data as usual. After training, the model is evaluated on the validation set, and the validation loss is printed. If the validation loss is lower than the best loss seen so far, best_loss is updated, wait is reset to zero, and the model’s parameters are saved. If not, wait is increased by one. If wait reaches the value of patience, the loop prints "Early stopping!" and breaks out of the training loop.

After training, the code checks if best_model_state is not None and then loads the best model weights back into the model. This ensures that you always use the model that performed best on the validation set, not just the one from the final epoch.

Here is an example of what the output might look like:

Epoch 0: Val Loss = 0.6931
Epoch 1: Val Loss = 0.6820
Epoch 2: Val Loss = 0.6755
Epoch 3: Val Loss = 0.6740
Epoch 4: Val Loss = 0.6742
Epoch 5: Val Loss = 0.6745
Epoch 6: Val Loss = 0.6747
Epoch 7: Val Loss = 0.6748
Epoch 8: Val Loss = 0.6749
Early stopping!

Epoch 0: Val Loss = 0.6931
Epoch 1: Val Loss = 0.6820
Epoch 2: Val Loss = 0.6755
Epoch 3: Val Loss = 0.6740
Epoch 4: Val Loss = 0.6742
Epoch 5: Val Loss = 0.6745
Epoch 6: Val Loss = 0.6747
Epoch 7: Val Loss = 0.6748
Epoch 8: Val Loss = 0.6749
Early stopping!

In this example, the validation loss improved for the first few epochs but then stopped getting better. After 5 epochs without improvement, training stopped early. The model parameters from the epoch with the lowest validation loss are restored at the end.

Introduction: Why Use Early Stopping?

How Early Stopping Works

Best Model vs. Last Model

Join the 1M+ learners on CodeSignal

Example: Adding Early Stopping and Saving the Best Model in PyTorch

Tuning Early Stopping

Summary And Practice Preview