Early stopping in PyTorch: patience and the best checkpoint

Early stopping in PyTorch ends training when a validation metric has gone too long without a meaningful improvement. You choose the metric, the improvement threshold, and the number of validation checks to wait. Keep the best model state separately so that stopping does not leave you with worse weights from a later check.

This guide builds a small CPU training loop with an explicit stopping rule. Sixteen checks passed with PyTorch 2.15.0 and Python 3.12.14, covering ties, threshold boundaries, invalid metrics, uneven validation batches, and checkpoint recovery. The synthetic example tests behavior; it does not establish better accuracy or faster convergence on a real dataset.

Define what patience counts

Call the stopping rule once after each complete validation pass. With one pass per epoch, three failed checks mean three epochs without a qualifying improvement. If validation runs every five epochs, the same setting covers three validation intervals. Calling the rule after every validation batch changes its meaning and makes the result depend on batching.

Use a fixed validation split for selection and keep the test set out of stopping and tuning decisions. Log the metric, validation-check number, patience counter, and selected checkpoint. Changes to the validation data or metric calculation can make consecutive results incomparable.

The helper below minimizes validation loss. It is code written for this guide, not a built-in PyTorch callback or a copy of Ignite’s behavior. If you use an engine, PyTorch-Ignite provides an EarlyStopping handler with its own options and event wiring. Check that library’s definitions before moving settings between implementations.

Separate patience from the best checkpoint

Our helper keeps two loss values. reference is the loss at the last patience reset. A new loss must be strictly below reference - min_delta to reset the counter. Smaller gains accumulate against that reference. best is the lowest loss observed, even when the improvement is too small to reset patience.

This separation preserves the actual lowest-loss weights. With patience=3 and min_delta=0.125, the following controlled sequence produces these results:

Check	Loss	Checks without reset	Decision
1	1.0000	0	Initialize reference and best weights
2	0.9375	1	Save lower loss; keep reference
3	0.8750	2	Exact threshold equality does not reset patience
4	0.8125	0	Save weights and reset patience
5	0.8125	1	Tie; keep the earlier weights
6	0.7500	2	Save lower loss; keep reference
7	0.7500	3	Stop and restore check 6

These values make the boundary easy to inspect, not a recommendation for your loss scale. In this implementation, patience must be at least one; patience one stops at the first check without a qualifying improvement. A zero threshold accepts any strictly lower loss. Ties never reset the counter.

Build the stopping helper

Save the following four Python blocks, in order, in a file named early_stopping.py. The first block defines the helper and imports used by the rest of the example:

import copy
import math

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


class EarlyStopper:
    """Minimize loss; count completed validation checks, not batches."""

    def __init__(self, patience=3, min_delta=0.0):
        if type(patience) is not int or patience < 1:
            raise ValueError("patience must be an integer >= 1")
        if not math.isfinite(min_delta) or min_delta < 0:
            raise ValueError("min_delta must be finite and >= 0")
        self.patience = patience
        self.min_delta = min_delta
        self.reference = math.inf
        self.best = math.inf
        self.best_state = None
        self.best_check = 0
        self.checks = 0
        self.bad_checks = 0
        self.stopped = False

    def step(self, loss, model):
        loss = float(loss)
        if not math.isfinite(loss):
            raise ValueError("validation loss must be finite")
        if self.stopped:
            raise RuntimeError("this run has already stopped")
        self.checks += 1
        if loss < self.best:
            self.best = loss
            self.best_check = self.checks
            self.best_state = copy.deepcopy(model.state_dict())
        if loss < self.reference - self.min_delta:
            self.reference = loss
            self.bad_checks = 0
        else:
            self.bad_checks += 1
        self.stopped = self.bad_checks >= self.patience
        return self.stopped

    def restore(self, model):
        if self.best_state is None:
            raise RuntimeError("no finite validation result was recorded")
        model.load_state_dict(self.best_state)

NaN or infinity raises an error before changing the monitor state. Treat a nonfinite validation result as a calculation or training failure to investigate, rather than silently waiting out patience. A stopped helper also refuses further checks; create a new instance for a new run.

The snapshot uses a deep copy because a plain model.state_dict() retains references to the model’s state. Later updates could otherwise change your saved “best” weights. PyTorch’s saving and loading tutorial documents this issue. This CPU example keeps a second model state in memory; account for that memory before adapting it to a large model.

Calculate one consistent validation loss

Append this validation function. It sums squared error across all target elements and divides once at the end. That matches the element-wise mean described by MSELoss, including an incomplete final batch.

@torch.no_grad()
def validation_mse(model, loader):
    model.eval()
    squared_error = 0.0
    elements = 0
    for x, y in loader:
        prediction = model(x)
        if prediction.shape != y.shape:
            raise ValueError("prediction and target shapes must match")
        squared_error += nn.functional.mse_loss(
            prediction, y, reduction="sum"
        ).item()
        elements += y.numel()
    if elements == 0:
        raise ValueError("validation loader is empty")
    return squared_error / elements

Averaging batch means without weighting them gives a short final batch too much influence. Our test used four zero-error elements and one element with squared error 100: the correct overall mean was 20. The function also rejects an empty loader and mismatched prediction/target shapes, preventing accidental broadcasting in this example. A masked or weighted loss needs its own denominator.

model.eval() selects evaluation behavior, while no_grad disables reverse-mode gradient recording for the validation computation. Neither replaces the other. The function leaves the model in evaluation mode; the training loop explicitly calls model.train() at the next epoch. Our guide to eval, no_grad, and inference_mode explains those choices.

Train, check, and restore

Append the CPU demonstration below. It fits a line to synthetic data, uses a separate validation grid, and validates once per epoch. The last validation batch has seven examples. No dataset download or GPU is needed.

def train_demo(max_epochs=200):
    torch.manual_seed(7)
    x = torch.linspace(-1, 1, 64).unsqueeze(1)
    train = DataLoader(TensorDataset(x, 2 * x + 0.5), batch_size=16)
    xv = torch.linspace(-0.95, 0.95, 23).unsqueeze(1)
    valid = DataLoader(TensorDataset(xv, 2 * xv + 0.5), batch_size=8)
    model = nn.Linear(1, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    stopper = EarlyStopper(patience=3, min_delta=0.001)
    history = []
    for epoch in range(1, max_epochs + 1):
        model.train()
        for xb, yb in train:
            optimizer.zero_grad(set_to_none=True)
            loss = nn.functional.mse_loss(model(xb), yb)
            loss.backward()
            optimizer.step()
        val_loss = validation_mse(model, valid)
        should_stop = stopper.step(val_loss, model)
        history.append(val_loss)
        print(f"epoch={epoch} loss={val_loss:.8f} "
              f"bad_checks={stopper.bad_checks}")
        if should_stop:
            break
    stopper.restore(model)
    restored_loss = validation_mse(model, valid)
    print(f"best_check={stopper.best_check} "
          f"restored_loss={restored_loss:.8f} stopped={stopper.stopped}")
    return model, stopper, history, restored_loss

The restore call sits after the loop, so it runs whether patience is exhausted or the epoch budget ends. In our tested run, training stopped at check 17 and the restored validation loss was approximately 0.00002780. The loss was still decreasing in small steps, which shows why a threshold can stop a run before every possible decrease is exhausted. These are control-flow settings, not tuned hyperparameters.

Save the selected weights and verify reloading

Append this final block. It saves the restored model, reloads it into the same architecture, and compares a prediction. The exclusive file mode refuses to overwrite an existing checkpoint.

if __name__ == "__main__":
    model, stopper, history, restored_loss = train_demo()
    # Refuse to overwrite an existing file. Choose a new run directory.
    with open("best-weights.pt", "xb") as checkpoint:
        torch.save(model.state_dict(), checkpoint)
    recovered = nn.Linear(1, 1)
    recovered.load_state_dict(torch.load(
        "best-weights.pt", map_location="cpu", weights_only=True
    ))
    recovered.eval()
    probe = torch.tensor([[0.25]])
    with torch.no_grad():
        torch.testing.assert_close(recovered(probe), model(probe))
    print("Saved weights reloaded; prediction matches.")

Run it from a new working directory in an environment with PyTorch installed:

python early_stopping.py

The script creates best-weights.pt in that directory. Use a different run directory for another execution. Our checks verified both the reload and refusal to replace an existing file. The saved weights support inference with the matching architecture; they do not contain a complete training-resume state.

For interruption recovery, save the optimizer, scheduler, progress, and relevant random-number state at a defined boundary. Preserve this helper’s reference, best loss and snapshot, counters, and settings too. Recreating it halfway through training resets patience. Our PyTorch checkpoint guide covers broader resume requirements; this example does not implement interruption recovery.

Coordinate stopping with the learning-rate schedule

If you use ReduceLROnPlateau, pass the completed validation metric to it at the intended interval. Choose a stopping budget that allows the reduced rate to be used and evaluated. A reduction immediately before breaking the loop cannot affect another training update.

The scheduler and this helper have different counter semantics. Here, three bad checks trigger stopping; ReduceLROnPlateau allows its configured number of bad checks before a subsequent one triggers reduction. Threshold modes can differ as well. Inspect both decisions in logs rather than copying the same patience value into each. Our learning-rate scheduler guide covers call timing and state restoration.

What to check before using the loop

Choose a metric whose direction matches your objective. This helper only minimizes a finite scalar loss; it has no accuracy-maximization mode. Set the threshold in the metric’s units, inspect validation noise, and choose patience in validation intervals. Retain a maximum training budget even when using early stopping.

The CPU tests covered threshold equality, ties, resets, nonfinite values, saved-state independence, model buffers, weighted validation, empty inputs, file recovery, and both loop exit paths. They did not test CUDA, mixed precision, distributed training, or real-data model quality. Distributed metric aggregation and coordinated termination require a separate implementation and test.