Sauvegarder un modèle PyTorch et reprendre son entraînement

Pour sauvegarder un modèle PyTorch destiné à l’inférence, enregistrez son state_dict, puis rechargez-le dans une architecture identique. Pour reprendre l’entraînement, conservez aussi l’état de la boucle : optimiseur, planificateur du taux d’apprentissage, compteur de progression et tout état qui détermine le prochain lot ou la prochaine mise à jour.

Un checkpoint utilisable permet de reprendre un entraînement interrompu sans tout recommencer. La création d’un fichier ne suffit pas : démarrez un nouveau processus, restaurez le checkpoint et comparez sa prochaine étape avec celle d’une exécution ininterrompue. Ce guide propose un petit exemple sur CPU, puis précise les adaptations nécessaires pour les GPU et les entraînements distribués.

Choisir entre les poids du modèle et un checkpoint d’entraînement

Le guide de sauvegarde et de chargement PyTorch recommande d’enregistrer le dictionnaire d’état du modèle plutôt que de sérialiser l’objet complet. Un dictionnaire de checkpoint peut également contenir l’état de l’optimiseur et la progression de l’entraînement.

Conservez l’architecture, la configuration, le prétraitement et les versions des dépendances avec les fichiers. Les poids n’indiquent pas comment segmenter une phrase en tokens, redimensionner une image ou interpréter une prédiction. Pour un adaptateur, notez aussi le modèle de base et sa révision ; notre guide du fine-tuning LoRA présente ce travail.

• Inférence : restaurez les paramètres entraînés et les buffers persistants, puis configurez l’évaluation.
• Reprise de l’entraînement : restaurez le modèle et les états nécessaires à la boucle.
• Nouvelle expérience : repartir des poids avec un nouvel optimiseur peut être volontaire ; consignez alors une nouvelle exécution.

Charger un checkpoint avec des paramètres explicites

Pour les fichiers composés de tenseurs et de dictionnaires de ce guide, utilisez torch.load(path, map_location="cpu", weights_only=True). L’API de chargement permet de choisir où restaurer le stockage des tenseurs. Le CPU convient notamment lorsque la machine de sauvegarde disposait d’un GPU absent de la machine de reprise.

weights_only limite la désérialisation ; cela ne signifie pas que le fichier peut contenir uniquement les poids du modèle. Notre exemple comprend aussi l’optimiseur, des tenseurs, des dictionnaires, des nombres et des chaînes. Les notes sur la sérialisation expliquent ce comportement et ses limites de sécurité. Chargez des fichiers de provenance fiable. Ne passez pas à weights_only=False uniquement pour faire disparaître une erreur inattendue.

Un exemple testé de sauvegarde et de reprise

L’exemple utilise un lot synthétique fixe, un petit réseau avec dropout, AdamW et un planificateur du taux d’apprentissage par étape. Il sauvegarde l’état après quatre mises à jour, puis calcule une cinquième mise à jour de référence. Un second processus restaure l’étape quatre et vérifie que sa cinquième mise à jour correspond à cette référence.

Enregistrez le code complet ci-dessous dans pytorch_checkpoint_demo.py. Exécutez-le dans un répertoire de travail distinct, avec PyTorch installé. Il crée un dossier checkpoint-demo et remplace ses deux fichiers de démonstration si vous le relancez. Cet exemple exact a été testé avec Python 3.12.14 et PyTorch 2.15.0+cpu sous Windows 11, sur un seul thread CPU. Aucun GPU n’est nécessaire.

import argparse
from pathlib import Path
import torch

CONFIG = {"inputs": 4, "hidden": 8, "outputs": 1, "dropout": 0.25}


def build(config):
    model = torch.nn.Sequential(
        torch.nn.Linear(config["inputs"], config["hidden"]),
        torch.nn.ReLU(),
        torch.nn.Dropout(config["dropout"]),
        torch.nn.Linear(config["hidden"], config["outputs"]),
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)
    return model, optimizer, scheduler


def train_step(model, optimizer, scheduler):
    # Fixed synthetic batch; no data-loader state is needed for this example.
    x = torch.arange(16, dtype=torch.float32).reshape(4, 4) / 10
    y = x.sum(dim=1, keepdim=True)
    model.train()
    optimizer.zero_grad(set_to_none=True)
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()
    scheduler.step()  # This example schedules once per optimizer update.
    return loss.item()


def snapshot(model, optimizer, scheduler, step, loss):
    return {
        "format_version": 1,
        "torch_version": str(torch.__version__),
        "model_config": CONFIG,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "completed_steps": step,
        "last_loss": loss,
        "torch_rng": torch.get_rng_state(),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=["save", "resume"])
    parser.add_argument("--directory", default="checkpoint-demo")
    args = parser.parse_args()
    folder = Path(args.directory)
    folder.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)

    if args.mode == "save":
        torch.manual_seed(7)
        model, optimizer, scheduler = build(CONFIG)
        for step in range(1, 5):
            loss = train_step(model, optimizer, scheduler)
        path = folder / "step-4.pt"
        torch.save(snapshot(model, optimizer, scheduler, step, loss), path)
        loss = train_step(model, optimizer, scheduler)
        torch.save(snapshot(model, optimizer, scheduler, 5, loss),
                   folder / "reference-step-5.pt")
        print(f"Saved {path}; reference step 5 loss: {loss:.8f}")
        return

    checkpoint = torch.load(folder / "step-4.pt", map_location="cpu", weights_only=True)
    if checkpoint["format_version"] != 1:
        raise ValueError("Unsupported checkpoint format")
    model, optimizer, scheduler = build(checkpoint["model_config"])
    model.load_state_dict(checkpoint["model"])
    scheduler.load_state_dict(checkpoint["scheduler"])
    optimizer.load_state_dict(checkpoint["optimizer"])
    # Restore RNG after constructing objects, immediately before the next step.
    torch.set_rng_state(checkpoint["torch_rng"])
    loss = train_step(model, optimizer, scheduler)
    step = checkpoint["completed_steps"] + 1
    reference = torch.load(folder / "reference-step-5.pt", map_location="cpu", weights_only=True)
    assert step == reference["completed_steps"]
    assert loss == reference["last_loss"]
    for key, value in model.state_dict().items():
        torch.testing.assert_close(value, reference["model"][key], rtol=0, atol=0)
    assert scheduler.state_dict() == reference["scheduler"]
    assert optimizer.param_groups[0]["lr"] == reference["optimizer"]["param_groups"][0]["lr"]
    print(f"Resume matched reference step {step}; loss: {loss:.8f}")


if __name__ == "__main__":
    main()

Exécuter les commandes dans deux processus distincts

python pytorch_checkpoint_demo.py save
python pytorch_checkpoint_demo.py resume

Lors du test consigné, le processus de reprise a affiché Resume matched reference step 5; loss: 10.83563042. Le script vérifie le compteur, la perte, les tenseurs du modèle, l’état du planificateur et le taux d’apprentissage. Cette perte correspond uniquement à cet exemple ; elle ne mesure ni la qualité d’un modèle utile ni les performances du matériel.

Des contrôles supplémentaires ont comparé l’état complet de l’optimiseur et testé des restaurations incomplètes. Omettre l’optimiseur modifiait le modèle après la mise à jour. Omettre l’état aléatoire modifiait la perte et le modèle. Omettre le planificateur laissait son compteur incorrect, même si la mise à jour suivante correspondait encore à la référence. Une seule perte identique ne prouve donc pas que la reprise est correcte.

Restaurer l’état d’entraînement dans le bon ordre

Créez le modèle, l’optimiseur et le planificateur avant de charger leurs états. Les groupes de paramètres de l’optimiseur doivent correspondre à ceux de l’exécution initiale. La documentation du chargement de l’optimiseur avertit qu’initialiser un planificateur après sa restauration peut écraser les taux d’apprentissage chargés. L’exemple initialise les deux objets, charge le planificateur, puis l’optimiseur.

Restaurez l’état aléatoire après la construction des objets, juste avant de continuer. La création des couches peut consommer des nombres aléatoires. Ici, le générateur CPU restauré contrôle le dropout ; le lot d’entrée reste fixe.

Définissez précisément le compteur. Ici, completed_steps=4 signifie que quatre mises à jour de l’optimiseur sont terminées : la suivante est la cinquième. Dans une boucle par époques, distinguez la dernière époque terminée de la prochaine à exécuter. Avec l’accumulation de gradients, privilégiez une sauvegarde après une mise à jour complète, sauf si votre code conserve aussi l’accumulation inachevée.

Sauvegarder et charger les poids pour l’inférence

Avec la même fonction de construction, l’exemple suivant extrait les poids du checkpoint d’entraînement et les charge dans un nouveau modèle. Placez-le à côté de pytorch_checkpoint_demo.py et exécutez-le après la commande de sauvegarde.

from pathlib import Path
import torch
from pytorch_checkpoint_demo import build

folder = Path("checkpoint-demo")
checkpoint = torch.load(folder / "step-4.pt", map_location="cpu", weights_only=True)
model, _, _ = build(checkpoint["model_config"])
model.load_state_dict(checkpoint["model"])
torch.save(model.state_dict(), folder / "inference-weights.pt")

loaded_model, _, _ = build(checkpoint["model_config"])
weights = torch.load(folder / "inference-weights.pt", map_location="cpu", weights_only=True)
loaded_model.load_state_dict(weights)
loaded_model.eval()
with torch.inference_mode():
    prediction = loaded_model(torch.ones(1, 4))
print(prediction)

model.eval() sélectionne le comportement d’évaluation des couches, dont le dropout. torch.inference_mode() agit sur le calcul des gradients ; il n’appelle pas eval() à votre place. La documentation du mode inférence précise cette distinction. Utilisez de nouveau model.train() pour reprendre l’entraînement.

Séparer le meilleur modèle du dernier checkpoint

Un modèle sélectionné selon la perte de validation et le dernier checkpoint ont des rôles différents. Le premier est le candidat à évaluer ou à déployer ; le second limite le travail à refaire après une interruption. Conservez les deux, avec le nom de la métrique, l’étape et l’identifiant de l’exécution.

Ne gardez pas un « meilleur » état uniquement avec best = model.state_dict() en mémoire pendant que l’entraînement continue. L’API Module précise que le dictionnaire contient des références aux paramètres et aux buffers. Enregistrez-le immédiatement ou créez une copie indépendante. Sinon, les mises à jour suivantes peuvent modifier l’état que vous pensiez conserver.

Adapter la reprise aux GPU, à la précision mixte et aux données mélangées

Définir le périmètre de reproductibilité

Consignez le matériel, les versions logicielles, l’ordre des données et la précision du test. PyTorch ne garantit pas des résultats identiques entre versions, plateformes ou exécutions CPU et GPU ; consultez ses recommandations sur la reproductibilité. Conserver une graine ne revient pas à sauvegarder l’état actuel de chaque générateur.

Un pipeline réel peut nécessiter les états aléatoires de Python, NumPy et CUDA, l’état de l’échantillonneur, celui du générateur du chargeur de données et la position dans le jeu de données. Les processus de chargement, les transformations aléatoires et le préchargement compliquent la reprise en milieu d’époque. Choisissez un point que votre pipeline peut restaurer et vérifiez le prochain lot autant que la prochaine mise à jour.

Restaurer l’état de précision mixte lorsqu’il est utilisé

Si l’entraînement utilise un ajusteur d’échelle des gradients, sauvegardez et restaurez son state_dict avec le modèle et l’optimiseur. La recette AMP officielle décrit les points de sauvegarde, notamment après scaler.update(). Cet exemple CPU ne teste ni AMP ni CUDA.

Conserver les checkpoints dans un stockage accessible au prochain job

Un checkpoint placé sur le disque local d’une instance qui disparaît ne peut pas aider le job de remplacement. Choisissez un stockage dont la persistance et les accès conviennent, puis testez la restauration depuis un autre processus ou une autre machine. Pour les grandes charges partagées, le guide des systèmes de fichiers HPC présente les choix de stockage.

Fixez la fréquence selon le travail que vous pouvez accepter de refaire et le coût mesuré d’une sauvegarde. Conservez des fichiers versionnés et au moins un checkpoint antérieur vérifié. Terminez l’écriture et la vérification avant de considérer le nouveau fichier comme la dernière copie utilisable. Avec un stockage distant, confirmez la fin du transfert. L’existence d’un nouveau nom de fichier ne justifie pas, à elle seule, la suppression du point de reprise précédent.

La démonstration utilise des écritures locales synchrones. Elle n’implémente pas une politique de stockage de production, qui doit aussi gérer les échecs, la conservation, la capacité et l’association entre fichiers, code du modèle et révision du jeu de données.

Choisir la méthode adaptée aux entraînements plus importants

Le checkpointing des activations et les checkpoints durables répondent à des besoins distincts. Le checkpointing des activations recalcule des résultats intermédiaires pendant la rétropropagation pour réduire l’utilisation de la mémoire. Il ne crée pas de fichier de reprise pour une autre machine.

Pour un modèle partitionné, utilisez une méthode prévue pour votre pile d’entraînement distribué. PyTorch Distributed Checkpoint gère des états répartis entre processus et impose ses propres conditions de compatibilité. Ne supposez pas que le fichier state_dict ordinaire d’un processus contient tout le modèle distribué.

Questions courantes sur la sauvegarde et le chargement

Pourquoi ne pas enregistrer l’objet modèle complet ?

torch.save(model, path) lie le fichier aux définitions des classes Python et aux chemins d’importation du modèle. Une réorganisation du code peut empêcher son chargement. Le dictionnaire d’état rend les paramètres enregistrés explicites, mais la reconstruction de l’architecture nécessite toujours un code compatible.

L’extension .pt ou .pth change-t-elle le contenu ?

Il s’agit de conventions de nommage. Vérifiez la structure enregistrée : le dictionnaire d’état d’un modèle et un checkpoint d’entraînement nécessitent des codes de chargement différents. Ajoutez une version à votre propre format de checkpoint pour signaler clairement les changements incompatibles.

Faut-il utiliser strict=False pour corriger une erreur ?

Vérifiez d’abord l’architecture, les dimensions des paramètres, les noms des clés et le type de fichier. Un chargement partiel peut servir une expérience volontaire d’apprentissage par transfert. Il ne remplace pas la restauration du modèle attendu. Examinez les clés manquantes ou inattendues.

Faut-il commencer un nouveau déploiement avec TorchScript ?

TorchScript est déprécié ; les notes de version PyTorch 2.14 le confirment. Choisissez une méthode adaptée au moteur d’exécution cible, par exemple un workflow pris en charge avec torch.export. L’export d’un artefact d’inférence est distinct de la sauvegarde d’un état permettant de reprendre l’entraînement.