Scheduler PyTorch : choisir et ajuster le taux d’apprentissage

Un scheduler de taux d’apprentissage PyTorch modifie le taux conservé par l’optimiseur. Avant d’en choisir un, définissez ce qui fait avancer le calendrier : une époque terminée, une mise à jour de l’optimiseur ou un résultat de validation. Déplacer le même appel dans la boucle change le déroulement de l’entraînement.

Ce guide compare quatre options courantes et propose trois petits exemples exécutables. Quatorze contrôles de fonctionnement ont réussi sur CPU avec PyTorch 2.15.0 et Python 3.12.14, dont une reprise depuis un checkpoint. Ils vérifient les comportements présentés, sans démontrer un gain de précision, de convergence ou de performance GPU.

Choisir le scheduler et son unité

Scheduler	Ce qui le pilote	Décision préalable
StepLR	Un intervalle fixe entre les réductions	La fréquence des appels à step ; ici, une fois par époque
CosineAnnealingLR	Une courbe cosinus sur une durée prévue	Le nombre d’appels représenté par T_max
ReduceLROnPlateau	Une métrique transmise à step	La métrique, son sens et la fréquence de validation
OneCycleLR	Un nombre défini de mises à jour	Le nombre de mises à jour de l’optimiseur pendant l’entraînement

Pour un calendrier par époque, appelez le scheduler après les mises à jour de cette époque. Pour un calendrier par mise à jour, appelez-le après la mise à jour correspondante. La documentation des optimiseurs précise cet ordre : appeler le scheduler avant l’optimiseur peut faire sauter le premier taux prévu.

Conservez cette unité dans la configuration et les journaux de l’expérience. Le scheduler ne compte pas automatiquement les données. Modifier la taille des lots, supprimer le dernier lot ou changer l’accumulation des gradients peut modifier le nombre de mises à jour. Notre guide DataLoader explique ces choix de découpage.

StepLR : distinguer le taux utilisé du prochain taux

StepLR multiplie le taux par gamma à chaque intervalle configuré. Enregistrez cet exemple autonome sous step_lr.py, puis exécutez-le dans un environnement où PyTorch est installé :

import torch
from torch import nn

x = torch.tensor([[0.0], [1.0], [2.0], [3.0]])
y = 2 * x + 1


def build():
    torch.manual_seed(7)
    model = nn.Linear(1, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer, step_size=2, gamma=0.5
    )
    return model, optimizer, scheduler


def train_epoch(model, optimizer):
    model.train()
    for start in range(0, len(x), 2):
        optimizer.zero_grad(set_to_none=True)
        loss = nn.functional.mse_loss(model(x[start:start + 2]), y[start:start + 2])
        loss.backward()
        optimizer.step()


if __name__ == "__main__":
    model, optimizer, scheduler = build()
    for epoch in range(6):
        used_lr = optimizer.param_groups[0]["lr"]
        train_epoch(model, optimizer)
        scheduler.step()
        next_lr = optimizer.param_groups[0]["lr"]
        print(f"epoch={epoch + 1} used={used_lr:.4f} next={next_lr:.4f}")

Chaque époque contient deux mises à jour de l’optimiseur. Le scheduler avance une fois après les deux : step_size=2 représente donc deux époques ici. Ce petit problème de régression fournit des gradients ; ses paramètres illustrent le mécanisme, sans constituer une recommandation pour un jeu de données réel.

L’exécution a produit les taux suivants :

Époque	Taux utilisé	Taux préparé pour l’époque suivante
1	0,1000	0,1000
2	0,1000	0,0500
3	0,0500	0,0500
4	0,0500	0,0250
5	0,0250	0,0250
6	0,0250	0,0125

Le dernier taux de 0,0125 a été préparé, mais jamais utilisé dans ces six époques. Nommer ces deux moments évite une erreur fréquente de journalisation. Avec plusieurs groupes de paramètres, consignez le lr de chaque groupe. L’API des schedulers propose aussi get_last_lr() pour les derniers taux calculés.

ReduceLROnPlateau : avancer après la validation

ReduceLROnPlateau reçoit une métrique. Utilisez mode="min" si une valeur basse est préférable, ou mode="max" si une valeur haute est préférable. Transmettez un résultat de validation calculé de façon cohérente, à un intervalle défini.

Ce script distinct utilise des métriques fictives pour rendre la patience observable. Il n’entraîne aucun modèle et ne mesure aucune qualité de validation :

import torch

parameter = torch.nn.Parameter(torch.tensor([0.0]))
optimizer = torch.optim.SGD([parameter], lr=0.1)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode="min", factor=0.5, patience=1,
    threshold=0.0, threshold_mode="abs"
)

# Synthetic metrics to inspect scheduler behavior; no model training here.
for validation_loss in [1.0, 0.8, 0.8, 0.8, 0.7]:
    scheduler.step(validation_loss)
    print(f"metric={validation_loss:.1f} next={optimizer.param_groups[0]['lr']:.3f}")

Les taux après les appels étaient 0.100, 0.100, 0.100, 0.050, 0.050. Avec un seuil absolu nul, les valeurs répétées de 0,8 n’améliorent pas la meilleure valeur. Une patience de 1 tolère un tel événement ; le deuxième déclenche la réduction. La dernière valeur, 0,7, améliore la métrique et remet à zéro le compteur d’événements sans amélioration.

Dans une boucle d’entraînement, terminez la validation avant scheduler.step(validation_loss). Utilisez la métrique agrégée prévue, pas le dernier lot d’entraînement rencontré. Si la fréquence de validation change, la même patience couvre une autre quantité d’entraînement. Choisissez le seuil, le cooldown et le taux minimal selon l’échelle et le bruit de la métrique.

OneCycleLR : compter les mises à jour effectives

OneCycleLR augmente puis réduit le taux sur un nombre prévu de mises à jour. Il accepte total_steps ou calcule ce total à partir des époques et des étapes par époque. Dans cet exemple autonome, chaque itération effectue une mise à jour de l’optimiseur :

import torch

parameter = torch.nn.Parameter(torch.tensor([1.0]))
optimizer = torch.optim.SGD([parameter], lr=0.1)
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=0.1, total_steps=8,
    pct_start=0.5, cycle_momentum=False
)

for update in range(8):
    optimizer.zero_grad(set_to_none=True)
    loss = parameter.square().sum()
    loss.backward()
    used_lr = optimizer.param_groups[0]["lr"]
    optimizer.step()
    scheduler.step()
    print(f"update={update + 1} used={used_lr:.8f}")

Les taux utilisés aux première, quatrième et huitième mises à jour étaient 0,004, 0,1 et 0,0000004. La construction du scheduler fixe le taux initial : le lr=0.1 donné auparavant à l’optimiseur n’est donc pas le premier taux utilisé. cycle_momentum=False concentre la démonstration sur le taux. Un neuvième appel au scheduler a provoqué une erreur, car les huit étapes prévues étaient épuisées.

Avec l’accumulation des gradients, plusieurs micro-lots peuvent produire une seule mise à jour. Comptez celle-ci une fois pour un calendrier par mise à jour et calculez le total à partir de la boucle réelle. Notre guide d’accumulation des gradients couvre cette limite et les derniers groupes incomplets.

La précision mixte ajoute un cas : les exemples AMP de PyTorch expliquent que le scaler peut sauter une mise à jour si les gradients contiennent des valeurs non finies. Alignez le scheduler par mise à jour sur les mises à jour réussies dans cette intégration. Nos tests CPU ne valident ni AMP ni l’entraînement distribué.

CosineAnnealingLR : définir la durée prévue

CosineAnnealingLR utilise T_max pour la durée du calendrier et eta_min pour le taux minimal. Son compteur avance à chaque appel. Un appel après chaque lot plutôt qu’après chaque époque donne au même T_max une autre durée d’entraînement.

Un contrôle CPU distinct, partant de 0,1 avec T_max=4 et eta_min=0.01, a atteint 0,01 après quatre appels. Ce scheduler n’implémente pas les warm restarts. Prévoyez la fin du calendrier sans supposer que des appels supplémentaires maintiendront indéfiniment le taux minimal.

Reprendre le calendrier avec l’état d’entraînement

Sauvegardez le modèle, l’optimiseur, l’état du scheduler et la progression à une limite clairement définie. Reconstruisez les objets correspondants avant de restaurer leur état. La référence du chargement de l’optimiseur prévient que créer le scheduler après la restauration de l’optimiseur peut écraser les taux restaurés.

Pour notre test StepLR, nous avons sauvegardé après la troisième époque, appel au scheduler compris. Après reconstruction des trois objets, nous avons restauré le modèle et le scheduler, puis l’optimiseur. Les taux restants étaient 0.05, 0.025, 0.025, comme dans l’exécution continue. Les paramètres du modèle, les buffers de momentum et l’état du scheduler étaient aussi strictement identiques dans cet exemple CPU déterministe.

Sans restauration du scheduler, les taux devenaient 0.05, 0.05, 0.025. Le premier était correct, mais la réduction suivante arrivait trop tard. Testez donc une reprise qui traverse une limite du calendrier. Notre guide des checkpoints PyTorch couvre les fichiers, l’état aléatoire et les autres besoins de reprise.

Vérifier le calendrier avant un entraînement long

• Réduction trop rapide : comptez les appels et vérifiez qu’un appel prévu par époque n’a pas été placé dans la boucle des lots.
• Aucune réduction sur plateau : examinez le sens de la métrique, le seuil, la patience, le cooldown et le taux minimal.
• Premier taux incorrect : inspectez l’optimiseur après la construction du scheduler et vérifiez l’ordre des appels.
• Reprise différente : comparez l’état du scheduler, les groupes de paramètres et la limite exacte de sauvegarde.

Conservez une trace des taux utilisés, des mises à jour terminées, des validations et de la configuration. Comparez les calendriers dans des conditions cohérentes d’entraînement et de validation. Une trace correcte vérifie que le calendrier s’est exécuté comme prévu ; son effet sur le modèle demande encore une expérience.