Arrêt anticipé dans PyTorch : patience et meilleur checkpoint

L’arrêt anticipé dans PyTorch met fin à l’entraînement lorsqu’une métrique de validation reste trop longtemps sans amélioration suffisante. Vous choisissez la métrique, le seuil d’amélioration et le nombre de validations à attendre. Conservez séparément le meilleur état du modèle pour éviter de garder des poids moins bons issus d’une validation ultérieure.

Ce guide construit une petite boucle CPU avec une règle d’arrêt explicite. Seize contrôles ont réussi avec PyTorch 2.15.0 et Python 3.12.14 : égalités, limites du seuil, métriques invalides, lots de validation inégaux et récupération du checkpoint. L’exemple synthétique vérifie le comportement du code, sans démontrer une meilleure précision ni une convergence plus rapide sur des données réelles.

Définir ce que compte la patience

Appelez la règle d’arrêt une fois après chaque passage complet de validation. Avec un passage par époque, trois validations sans progrès suffisant correspondent à trois époques. Si vous validez toutes les cinq époques, le même réglage couvre trois intervalles de validation. Appeler la règle après chaque lot de validation change son sens et rend le résultat dépendant du découpage en lots.

Utilisez un jeu de validation fixe pour la sélection et gardez le jeu de test en dehors des décisions d’arrêt et de réglage. Consignez la métrique, le numéro de validation, le compteur de patience et le checkpoint sélectionné. Modifier les données ou le calcul de la métrique peut rendre les résultats successifs incomparables.

Le helper ci-dessous minimise la perte de validation. Il a été écrit pour ce guide : ce n’est ni un callback intégré à PyTorch ni une reproduction du comportement d’Ignite. Si vous utilisez un moteur d’entraînement, PyTorch-Ignite propose un gestionnaire EarlyStopping avec ses propres options et événements. Vérifiez les définitions de la bibliothèque avant de transférer des réglages d’une implémentation à l’autre.

Séparer la patience du meilleur checkpoint

Notre helper conserve deux pertes. reference est la valeur au dernier rétablissement de la patience. Une nouvelle perte doit être strictement inférieure à reference - min_delta pour remettre le compteur à zéro. Les petits gains s’accumulent par rapport à cette référence. best est la plus faible perte observée, même si le gain est insuffisant pour réinitialiser la patience.

Cette séparation conserve les poids correspondant à la perte réellement la plus basse. Avec patience=3 et min_delta=0.125, cette séquence contrôlée donne les résultats suivants :

Validation	Perte	Validations sans remise à zéro	Décision
1	1,0000	0	Initialiser la référence et les meilleurs poids
2	0,9375	1	Conserver les nouveaux poids ; garder la référence
3	0,8750	2	L’égalité exacte au seuil ne réinitialise pas la patience
4	0,8125	0	Conserver les poids et réinitialiser la patience
5	0,8125	1	Égalité ; garder les poids précédents
6	0,7500	2	Conserver les nouveaux poids ; garder la référence
7	0,7500	3	Arrêter et restaurer la validation 6

Ces valeurs rendent la limite facile à examiner ; elles ne recommandent aucun seuil pour votre perte. Ici, la patience doit être au moins égale à un. Une patience de un arrête l’entraînement à la première validation sans amélioration suffisante. Un seuil nul accepte toute perte strictement inférieure. Une égalité ne remet jamais le compteur à zéro.

Construire le helper d’arrêt

Enregistrez les quatre blocs Python suivants, dans l’ordre, dans un fichier nommé early_stopping.py. Le premier définit le helper et les imports nécessaires à la suite :

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)

Une valeur NaN ou infinie déclenche une erreur avant toute modification de l’état de suivi. Recherchez la cause d’un résultat de validation non fini dans le calcul ou l’entraînement, plutôt que d’épuiser silencieusement la patience. Un helper déjà arrêté refuse aussi de nouvelles validations ; créez une instance pour chaque nouvel entraînement.

La sauvegarde utilise une copie profonde, car un simple model.state_dict() conserve des références à l’état du modèle. Les mises à jour suivantes pourraient modifier les « meilleurs » poids conservés. Le tutoriel de sauvegarde et de chargement PyTorch explique ce point. Cet exemple CPU garde un second état du modèle en mémoire ; prévoyez ce coût avant de l’adapter à un grand modèle.

Calculer une perte de validation cohérente

Ajoutez cette fonction. Elle additionne les erreurs quadratiques sur tous les éléments cibles, puis divise une seule fois à la fin. Cela correspond à la moyenne par élément décrite par MSELoss, même lorsque le dernier lot est incomplet.

@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

Faire la moyenne des pertes moyennes des lots sans les pondérer donne trop de poids au dernier petit lot. Notre test comptait quatre éléments sans erreur et un élément dont l’erreur quadratique valait 100 : la moyenne globale correcte était 20. La fonction rejette aussi un loader vide et des formes différentes entre prédictions et cibles, évitant ici un broadcasting accidentel. Une perte masquée ou pondérée demande son propre dénominateur.

model.eval() sélectionne le comportement d’évaluation, tandis que no_grad désactive l’enregistrement des gradients en mode inverse pour le calcul de validation. L’un ne remplace pas l’autre. La fonction laisse le modèle en mode évaluation ; la boucle appelle explicitement model.train() à l’époque suivante. Notre guide sur eval, no_grad et inference_mode (en anglais) détaille ces choix.

Entraîner, vérifier et restaurer

Ajoutez cette démonstration CPU. Elle ajuste une droite à des données synthétiques, utilise une grille de validation distincte et valide une fois par époque. Le dernier lot de validation contient sept exemples. Aucun téléchargement de données ni GPU n’est nécessaire.

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

L’appel de restauration se situe après la boucle : il s’exécute aussi bien après épuisement de la patience qu’à la fin du nombre d’époques autorisé. Lors de notre test, l’entraînement s’est arrêté à la validation 17, avec une perte restaurée d’environ 0.00002780. La perte diminuait encore par petits pas : un seuil peut donc arrêter l’entraînement avant l’épuisement de toutes les diminutions possibles. Ces valeurs servent à démontrer le déroulement du code ; ce ne sont pas des hyperparamètres optimisés.

Sauvegarder les poids retenus et vérifier leur rechargement

Ajoutez ce dernier bloc. Il sauvegarde le modèle restauré, le recharge dans la même architecture et compare une prédiction. Le mode exclusif refuse d’écraser un checkpoint existant.

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.")

Exécutez le fichier depuis un nouveau répertoire de travail, dans un environnement où PyTorch est installé :

python early_stopping.py

Le script crée best-weights.pt dans ce répertoire. Utilisez un autre répertoire pour une nouvelle exécution. Nos contrôles ont vérifié le rechargement et le refus de remplacer un fichier existant. Ces poids permettent l’inférence avec l’architecture correspondante ; ils ne constituent pas un état complet de reprise d’entraînement.

Pour reprendre après interruption, sauvegardez l’optimiseur, le scheduler, la progression et les états aléatoires utiles à une limite définie. Conservez aussi la référence du helper, sa meilleure perte et les poids associés, ses compteurs et ses réglages. Recréer le helper en cours d’entraînement réinitialise la patience. Notre guide des checkpoints PyTorch traite les besoins de reprise ; cet exemple ne met pas en œuvre la récupération après interruption.

Coordonner l’arrêt avec le taux d’apprentissage

Si vous utilisez ReduceLROnPlateau, transmettez-lui la métrique du passage de validation terminé à l’intervalle prévu. Choisissez un délai d’arrêt qui laisse le temps d’utiliser le taux réduit, puis d’en évaluer l’effet. Une réduction juste avant de sortir de la boucle ne peut agir sur aucune nouvelle mise à jour.

Le scheduler et ce helper interprètent leurs compteurs différemment. Ici, trois validations insuffisantes déclenchent l’arrêt ; ReduceLROnPlateau tolère le nombre configuré avant qu’une validation supplémentaire ne déclenche la réduction. Les modes de seuil peuvent aussi différer. Examinez les deux décisions dans les journaux plutôt que de recopier la même patience. Notre guide des schedulers (en anglais) explique l’ordre des appels et la restauration de leur état.

Vérifier la boucle avant de l’utiliser

Choisissez une métrique dont le sens correspond à votre objectif. Ce helper minimise uniquement une perte scalaire finie ; il ne propose pas de mode de maximisation de la précision. Réglez le seuil dans les unités de la métrique, examinez le bruit de validation et choisissez la patience en intervalles de validation. Conservez un nombre maximal d’époques même avec un arrêt anticipé.

Les tests CPU ont couvert l’égalité au seuil, les ex æquo, les remises à zéro, les valeurs non finies, l’indépendance de l’état conservé, les buffers du modèle, la validation pondérée, les entrées vides, la récupération depuis un fichier et les deux sorties de boucle. Ils n’ont testé ni CUDA, ni la précision mixte, ni l’entraînement distribué, ni la qualité sur des données réelles. L’agrégation distribuée de la métrique et l’arrêt coordonné exigent une implémentation et des tests distincts.