Partage de technologie

Comment continuer la formation après avoir enregistré le modèle pytorch localement

2024-07-12

한어Русский языкEnglishFrançaisIndonesianSanskrit日本語DeutschPortuguêsΕλληνικάespañolItalianoSuomalainenLatina

Dans PyTorch, vous pouvez enregistrer et charger le modèle et continuer la formation en suivant ces étapes :

  1. Enregistrer le modèle

    Il existe généralement deux manières d'enregistrer un modèle :

    • Enregistrez l'intégralité du modèle (y compris la structure du réseau, les poids, etc.) :

      torch.save(model, 'model.pth')
    • Seul le state_dict du modèle (contenant uniquement les paramètres de poids) est enregistré. Cette méthode est recommandée car elle économise de l'espace de stockage et est plus flexible lors du chargement :

      torch.save(model.state_dict(), 'model_weights.pth')
  2. Charger le modèle

    En conséquence, il existe deux manières de charger le modèle :

    • Si vous avez déjà enregistré l'intégralité du modèle, vous pouvez le charger directement de la manière suivante :

      model = torch.load('model.pth')
    • Si vous avez uniquement enregistré state_dict auparavant, vous devez instancier un modèle avec la même structure que le modèle d'origine, puis transmettreload_state_dict()Méthode pour charger des poids :

      1. # 实例化一个与原模型结构相同的模型
      2. model = YourModelClass()
      3. # 加载保存的state_dict
      4. model.load_state_dict(torch.load('model_weights.pth'))
      5. # 确保将模型转移到正确的设备上(例如GPU或CPU)
      6. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
      7. model.to(device)
  3. continuer la formation

    Après avoir chargé le modèle, vous pouvez continuer la formation. Assurez-vous d'avoir défini la fonction de perte et l'optimiseur et que leurs états sont correctement chargés (si vous les avez enregistrés précédemment).Ensuite, suivez simplement le processus de formation normal

    1. # 定义损失函数和优化器
    2. criterion = nn.CrossEntropyLoss()
    3. optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    4. # 如果之前保存了优化器状态,也可以加载
    5. optimizer.load_state_dict(torch.load('optimizer.pth'))
    6. # 开始训练
    7. for epoch in range(num_epochs):
    8. for inputs, labels in dataloader:
    9. inputs, labels = inputs.to(device), labels.to(device)
    10. optimizer.zero_grad()
    11. outputs = model(inputs)
    12. loss = criterion(outputs, labels)
    13. loss.backward()
    14. optimizer.step()

De cette façon, vous pouvez continuer à entraîner le modèle à partir de l'endroit où vous l'avez enregistré pour la dernière fois.