La meilleure façon de sauvegarder un modèle formé dans PyTorch ?

Question

Je cherchais des moyens alternatifs pour sauvegarder un modèle formé dans PyTorch. Jusqu'à présent, j'ai trouvé deux alternatives.

1. torche.save() pour sauvegarder un modèle et torche.load() pour charger un modèle.
2. modèle.state_dict() pour sauvegarder un modèle formé et modèle.load_state_dict() pour charger le modèle sauvegardé.

Je suis tombé sur ceci discussion où l'approche 2 est recommandée par rapport à l'approche 1.

Ma question est la suivante : pourquoi la deuxième approche est-elle préférable ? Est-ce seulement parce que torche.nn les modules ont ces deux fonctions et nous sommes encouragés à les utiliser ?

Answer 1

Sauvegarde locale

La façon dont vous enregistrez votre modèle dépend de la façon dont vous voulez y accéder à l'avenir. Si vous pouvez appeler une nouvelle instance de la classe model tout ce que vous avez à faire est de sauvegarder/charger les poids du modèle avec model.state_dict() :

# Save:
torch.save(old_model.state_dict(), PATH)

# Load:
new_model = TheModelClass(*args, **kwargs)
new_model.load_state_dict(torch.load(PATH))

Si vous ne pouvez pas le faire pour une raison quelconque (ou si vous préférez une syntaxe plus simple), vous pouvez sauvegarder le modèle entier (en fait une référence au(x) fichier(s) définissant le modèle, ainsi que son state_dict) avec torch.save() :

# Save:
torch.save(old_model, PATH)

# Load:
new_model = torch.load(PATH)

Mais comme il s'agit d'une référence à l'emplacement des fichiers définissant la classe modèle, ce code n'est pas portable à moins que ces fichiers ne soient également portés dans la même structure de répertoire.

Sauvegarde dans le nuage - TorchHub

Si vous souhaitez que votre modèle soit portable, vous pouvez facilement permettre qu'il soit importé à l'aide de la fonction torch.hub . Si vous ajoutez un hubconf.py dans un répertoire github, qui peut être facilement appelé à partir de PyTorch pour permettre aux utilisateurs de charger votre modèle avec ou sans poids :

hubconf.py ( github.com/repo_owner/repo_name )

dependencies = ['torch']
from my_module import mymodel as _mymodel

def mymodel(pretrained=False, **kwargs):
    return _mymodel(pretrained=pretrained, **kwargs)

Modèle de chargement :

new_model = torch.hub.load('repo_owner/repo_name', 'mymodel')
new_model_pretrained = torch.hub.load('repo_owner/repo_name', 'mymodel', pretrained=True)

Answer 2

De nos jours, tout est écrit dans le tutoriel officiel : https://pytorch.org/tutorials/beginner/saving_loading_models.html

Vous avez plusieurs options sur la façon de sauvegarder et sur ce que vous voulez sauvegarder et tout est expliqué dans ce tutoriel.