Wie man Modelle in PyTorch speichert und lädt

Einführung

Sektionen

Verwenden Sie state_dict (empfohlen)

Speichern

torch.save(model.state_dict(), 'save/to/path/model.pth')

Laden

model = MyModelDefinition(args)
model.load_state_dict(torch.load('load/from/path/model.pth'))

Vorteile:

PyTorch stützt sich intern auf das Pickle-Modul von Python. Das Python-Verzeichnis kann leicht gepickelt, entpickelt, aktualisiert und wiederhergestellt werden. Das Speichern von Modellen mit state_dict bietet daher mehr Flexibilität.
Sie können auch den Optimierungsstatus, die Hyperparameter usw. als Schlüssel-Wert-Paare zusammen mit dem state_dict des Modells speichern. Nach der Wiederherstellung können Sie auf diese Paare wie auf ein gewöhnliches Python-Verzeichnis zugreifen. Wie das geht, werden wir im nächsten Abschnitt sehen.

Nachteile:

Sie benötigen die Modelldefinition, um das state_dict zu laden.

Tücken:

Stellen Sie sicher, dass Sie model.eval() aufrufen, wenn Sie eine Inferenz machen wollen.
Speichern Sie das Modell mit der Erweiterung .pt oder .pth.

Speichern und Laden des gesamten Modells

Speichern

torch.save(model, 'save/to/path/model.pt')

Laden

model = torch.load('load/from/path/model.pt')

Vorteile:

Der einfachste Weg, das gesamte Modell mit dem geringsten Codeaufwand zu speichern.
Die API zum Speichern und Laden ist intuitiver.

Nachteile:

Da das Pickle-Modul von Python intern verwendet wird, sind die serialisierten Daten an die spezifischen Klassen gebunden und die genaue Verzeichnisstruktur wird beim Speichern des Modells verwendet. Pickle speichert einfach einen Pfad zu der Datei, welche die jeweilige Klasse enthält. Dies wird beim Laden verwendet.
Sie können sich vorstellen, dass der Code nach dem Refactoring nicht mehr funktioniert, da das gespeicherte Modell möglicherweise nicht mit demselben Pfad verknüpft ist. Die Verwendung eines solchen Modells in einem anderen Projekt ist ebenfalls schwierig, da die Pfadstruktur beibehalten werden muss.

Tücken:

Stellen Sie sicher, dass Sie model.eval() aufrufen, wenn Sie Inferenzen durchführen wollen.
Speichern Sie das Modell mit der Erweiterung .pt oder .pth.

Speichern und Laden des Modells von einem Checkpoint aus

Speichern

torch.save({
            'epoch': EPOCH,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS,
            }, 'save/to/path/model.pth')

Laden

model = MyModelDefinition(args)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load('load/from/path/model.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

Tücken:

Um dies für das Training zu verwenden, rufen Sie model.train() auf.
Um dies für die Inferenz zu verwenden, rufen Sie model.eval() auf.

Verwendung von W&B-Artefakten für die Modellversionskontrolle



Als Artefakte speichern

# Import
import wandb
# Save your model.
torch.save(model.state_dict(), 'save/to/path/model.pth')
# Save as artifact for version control.
run = wandb.init(project='your-project-name')
artifact = wandb.Artifact('model', type='model')
artifact.add_file('save/to/path/model.pth')
run.log_artifact(artifact)
run.join()

Als Artefakte laden

import wandb
run = wandb.init()

artifact = run.use_artifact('entity/your-project-name/model:v0', type='model')
artifact_dir = artifact.download()

run.join()

Abbildung 1: Sehen Sie sich hier das Live-Artefakt-Dashboard an

Weights & Biases

Empfohlene Lektüre für alle, die sich für PyTorch interessieren