Sauvegarde et chargement de modèle PyTorch

Dans les projets d'apprentissage profond, la sauvegarde et le chargement du modèle constituent une étape cruciale. Les principales raisons incluent :

  1. Récupération après interruption de l'entraînement: Vous pouvez reprendre l'entraînement à partir d'un checkpoint lorsque le processus d'entraînement est interrompu de manière inattendue
  2. Déploiement du modèle: Déployez le modèle entraîné dans un environnement de production
  3. Partage de modèle: Pratique pour partager les résultats du modèle entre les membres de l'équipe
  4. Apprentissage par transfert: Sauvegardez les modèles pré-entraînés pour d'autres tâches
  5. Évaluation des performances: Sauvegardez les modèles issus de différentes étapes d'entraînement pour comparaison

Méthodes de base de sauvegarde et de chargement

Sauvegarde du modèle complet

Il s'agit de la méthode la plus simple, qui sauvegarde l'architecture et les paramètres du modèle :

Exemple

import torch
import torchvision.models as models

# Créer et entraîner un modèle
model = models.resnet18(pretrained=True)
# ... code d'entraînement ...

# Enregistrer l'intégralité du modèle
torch.save(model, 'model.pth')

# Charger l'intégralité du modèle
loaded_model = torch.load('model.pth')

Avantages：

  • Code simple et intuitif
  • Enregistre la structure complète du modèle

Inconvénients：

  • Taille de fichier plus importante
  • Dépend de la définition de la classe du modèle

Enregistrement des paramètres du modèle uniquement (recommandé)

Une méthode plus recommandée consiste à enregistrer uniquement le state_dict du modèle :

Exemple

# Enregistrer les paramètres du modèle
torch.save(model.state_dict(), 'model_weights.pth')

# Charger les paramètres du modèle
model = models.resnet18()  # Il faut d'abord créer un modèle avec la même architecture
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # Passer en mode évaluation

Avantages：

  • Taille de fichier plus petite
  • Plus flexible, peut être chargé dans différentes architectures
  • Meilleure compatibilité

Enregistrement et chargement de l'état d'entraînement

Dans les projets réels, nous devons généralement aussi sauvegarder l'état de l'optimiseur, l'époque et d'autres informations :

Exemple

# Enregistrer le point de contrôle
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    # Vous pouvez ajouter d'autres informations qui doivent être sauvegardées
}

torch.save(checkpoint, 'checkpoint.pth')

# Charger le point de contrôle
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()  # Ou model.train() selon vos besoins

Chargement de modèles sur différents appareils

Gestion de la compatibilité CPU/GPU

Exemple

# Spécifier map_location lors de l'enregistrement
torch.save(model.state_dict(), 'model_weights.pth')

# Charger sur le CPU (lorsque le modèle a été entraîné sur GPU)
device = torch.device('cpu')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))

# Charger sur GPU
device = torch.device('cuda')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))
model.to(device)

Chargement des modèles entraînés avec plusieurs GPU

Exemple

# Enregistrer un modèle multi-GPU
torch.save(model.module.state_dict(), 'multigpu_model.pth')

# Charger sur un seul GPU
model = ModelClass()
model.load_state_dict(torch.load('multigpu_model.pth'))

Conversion de modèles et compatibilité

Compatibilité des versions de PyTorch

Exemple

# Spécifiez _use_new_zipfile_serialization=True lors de l'enregistrement pour une meilleure compatibilité
torch.save(model.state_dict(), 'model.pth', _use_new_zipfile_serialization=True)

Conversion vers TorchScript

Exemple

# Convertir le modèle au format TorchScript
scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'model_scripted.pt')

# Charger le modèle TorchScript
loaded_script = torch.jit.load('model_scripted.pt')

Bonnes pratiques et problèmes courants

Bonnes pratiques

  1. Conventions de nommage: Utilisez des noms de fichiers explicites, tels queresnet18_epoch50.pth
  2. Sauvegarde périodique: Enregistrer un point de contrôle toutes les quelques époques
  3. Vérifier le chargement: Tester la fonctionnalité de chargement immédiatement après la sauvegarde
  4. Documentation: Enregistrer l'architecture du modèle et les paramètres d'entraînement
  5. Contrôle de version: Inclure les fichiers du modèle dans le système de contrôle de version

Solutions aux problèmes courants

Problème 1：Clé(s) manquante(s) dans state_dict
Solution: S'assurer que l'architecture du modèle correspond exactement, ou utiliserstrict=Falseparamètre :

model.load_state_dict(torch.load('model.pth'), strict=False)

Problème 2：Mémoire CUDA insuffisante
Solution: Placez-le d'abord sur le CPU lors du chargement :

Exemple

model.load_state_dict(torch.load('model.pth', map_location='cpu'))

Problème 3: Impossible de charger les modèles d'anciennes versions
Solution: Essayez de charger dans différentes versions de PyTorch, ou convertissez le format du modèle

Exemples d'applications pratiques

Processus de sauvegarde et de chargement du modèle de classification d'images

Exemple de code complet

Exemple

import torch
import torch.nn as nn
import torch.optim as optim

# Définir un modèle simple
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2)
   
    def forward(self, x):
        return self.fc(x)

# Initialisation
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# Simuler le processus d'entraînement
for epoch in range(5):
    # Simuler les étapes d'entraînement
    inputs = torch.randn(32, 10)
    labels = torch.randint(0, 2, (32,))
   
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
   
    # Enregistrer un point de contrôle toutes les 2 époques
    if epoch % 2 == 0:
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss.item(),
        }
        torch.save(checkpoint, f'checkpoint_epoch{epoch}.pth')
        print(f'Point de contrôle enregistré à l’époque {epoch}')

# Sauvegarde finale
torch.save(model.state_dict(), 'final_model.pth')

# Exemple de chargement
loaded_model = SimpleModel()
loaded_model.load_state_dict(torch.load('final_model.pth'))
loaded_model.eval()

# Tester le modèle chargé
test_input = torch.randn(1, 10)
with torch.no_grad():
    output = loaded_model(test_input)
print(f'Test output: {output}')
Autres extensions