Wie speichere ich ein trainiertes Modell in PyTorch?

Wie speichere ich ein trainiertes Modell in PyTorch? Das habe ich gelesen:

1. torch.save()/torch.load() dient zum Speichern/Laden eines serialisierbaren Objekts.
2. model.state_dict()/model.load_state_dict() dient zum Speichern/Laden des Modellzustands.

Gefunden diese Seite auf ihrem github repo:

Empfohlene Vorgehensweise zum Speichern eines Modells

Es gibt zwei Hauptansätze zum Serialisieren und Wiederherstellen eines Modells.

Die erste (empfohlene) speichert und lädt nur die Modellparameter:

torch.save(the_model.state_dict(), PATH)

Dann später:

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

Die zweite speichert und lädt das gesamte Modell:

torch.save(the_model, PATH)

Dann später:

the_model = torch.load(PATH)

In diesem Fall sind die serialisierten Daten jedoch an die spezifischen Klassen und die genaue verwendete Verzeichnisstruktur gebunden, sodass sie bei der Verwendung in anderen Projekten oder nach einigen ernsthaften Umgestaltungen auf verschiedene Weise beschädigt werden können.

Siehe auch: Speichern und laden Sie das Modell Abschnitt aus den offiziellen PyTorch-Tutorials.

Es hängt davon ab, was Sie tun möchten.

Fall Nr. 1: Speichern Sie das Modell, um es selbst für die Inferenz zu verwenden: Sie speichern das Modell, stellen es wieder her und ändern das Modell dann in den Evaluierungsmodus. Dies geschieht, weil Sie normalerweise haben BatchNorm und Dropout Layer, die sich standardmäßig im Zugmodus bei der Konstruktion befinden:

torch.save(model.state_dict(), filepath)

#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()

Fall Nr. 2: Speichern Sie das Modell, um das Training später fortzusetzen: Wenn Sie das zu speichernde Modell weiter trainieren müssen, müssen Sie mehr als nur das Modell speichern. Sie müssen auch den Zustand des Optimierers, Epochen, Punktzahl usw. speichern. Sie würden es so machen:

state = {
    'epoch': epoch,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    ...
}
torch.save(state, filepath)

Um das Training wieder aufzunehmen, würden Sie Folgendes tun: state = torch.load(filepath)und dann, um den Zustand jedes einzelnen Objekts wiederherzustellen, etwa so:

model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])

Da Sie das Training wieder aufnehmen, UNTERLASSEN SIE Anruf model.eval() sobald Sie die Zustände beim Laden wiederherstellen.

Fall Nr. 3: Modell, das von jemand anderem ohne Zugriff auf Ihren Code verwendet werden soll: In Tensorflow können Sie eine erstellen .pb Datei, die sowohl die Architektur als auch die Gewichtungen des Modells definiert. Dies ist sehr praktisch, besonders bei der Verwendung Tensorflow serve. Der äquivalente Weg, dies in Pytorch zu tun, wäre:

torch.save(model, filepath)

# Then later:
model = torch.load(filepath)

Dieser Weg ist immer noch nicht kugelsicher und da Pytorch immer noch viele Änderungen durchmacht, würde ich es nicht empfehlen.

Das Essiggurke Die Python-Bibliothek implementiert binäre Protokolle zum Serialisieren und Deserialisieren eines Python-Objekts.

Wenn du import torch (oder wenn Sie PyTorch verwenden) wird es import pickle für dich und du brauchst nicht anzurufen pickle.dump() und pickle.load() direkt, das sind die Methoden zum Speichern und Laden des Objekts.

In der Tat, torch.save() und torch.load() wird wickeln pickle.dump() und pickle.load() für dich.

EIN state_dict Die andere erwähnte Antwort verdient nur ein paar weitere Anmerkungen.

Was state_dict haben wir in PyTorch? Es sind eigentlich zwei state_dicts.

Das PyTorch-Modell ist torch.nn.Module was hat model.parameters() aufrufen, um lernbare Parameter (w und b) zu erhalten. Diese lernbaren Parameter, die einmal zufällig festgelegt wurden, werden im Laufe der Zeit aktualisiert, wenn wir lernen. Lernbare Parameter sind die ersten state_dict.

Der Zweite state_dict ist der Zustand des Optimierers dict. Sie erinnern sich, dass der Optimierer verwendet wird, um unsere lernbaren Parameter zu verbessern. Aber der Optimierer state_dict Ist repariert. Da gibt es nichts zu lernen.

Da state_dict Objekte sind Python-Wörterbücher, sie können einfach gespeichert, aktualisiert, geändert und wiederhergestellt werden, was den PyTorch-Modellen und -Optimierern eine große Modularität verleiht.

Lassen Sie uns ein supereinfaches Modell erstellen, um dies zu erklären:

import torch
import torch.optim as optim

model = torch.nn.Linear(5, 2)

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

print("Model weight:")    
print(model.weight)

print("Model bias:")    
print(model.bias)

print("---")
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

Dieser Code gibt Folgendes aus:

Model's state_dict:
weight      torch.Size([2, 5])
bias      torch.Size([2])
Model weight:
Parameter containing:
tensor([[ 0.1328,  0.1360,  0.1553, -0.1838, -0.0316],
        [ 0.0479,  0.1760,  0.1712,  0.2244,  0.1408]], requires_grad=True)
Model bias:
Parameter containing:
tensor([ 0.4112, -0.0733], requires_grad=True)
---
Optimizer's state_dict:
state      {}
param_groups      [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [140695321443856, 140695321443928]}]

Beachten Sie, dass dies ein Minimalmodell ist. Sie können versuchen, einen Stapel von sequentiellen hinzuzufügen

model = torch.nn.Sequential(
          torch.nn.Linear(D_in, H),
          torch.nn.Conv2d(A, B, C)
          torch.nn.Linear(H, D_out),
        )

Beachten Sie, dass nur Schichten mit lernbaren Parametern (Faltungsschichten, lineare Schichten usw.) und registrierte Puffer (Batchnorm-Schichten) Einträge in den Modellen haben state_dict.

Nicht lernbare Dinge gehören zum Optimizer-Objekt state_dictdie Informationen über den Zustand des Optimierers sowie die verwendeten Hyperparameter enthält.

Der Rest der Geschichte ist derselbe; in der Inferenzphase (dies ist eine Phase, in der wir das Modell nach dem Training verwenden) zum Vorhersagen; wir prognostizieren basierend auf den Parametern, die wir gelernt haben. Für die Inferenz müssen wir also nur die Parameter speichern model.state_dict().

torch.save(model.state_dict(), filepath)

Und später zu verwenden model.load_state_dict(torch.load(filepath)) model.eval()

Hinweis: Vergessen Sie nicht die letzte Zeile model.eval() Dies ist nach dem Laden des Modells entscheidend.

Versuchen Sie auch nicht zu speichern torch.save(model.parameters(), filepath). Das model.parameters() ist nur das Generatorobjekt.

Auf der anderen Seite, torch.save(model, filepath) speichert das Modellobjekt selbst, aber denken Sie daran, dass das Modell nicht über den Optimierer verfügt state_dict. Überprüfen Sie die andere ausgezeichnete Antwort von @ Jadiel de Armas, um das Zustandsdikt des Optimierers zu speichern.

Eine gängige PyTorch-Konvention besteht darin, Modelle entweder mit der Dateierweiterung .pt oder .pth zu speichern.

Gesamtes Modell speichern/laden

Speichern:

path = "username/directory/lstmmodelgpu.pth"
torch.save(trainer, path)

Belastung:

(Modellklasse muss irgendwo definiert werden)

model.load_state_dict(torch.load(PATH))
model.eval()

Wenn Sie das Modell speichern und das Training später fortsetzen möchten:

Einzel-GPU:
Speichern:

state = {
        'epoch': epoch,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
}
savepath="checkpoint.t7"
torch.save(state,savepath)

Belastung:

checkpoint = torch.load('checkpoint.t7')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']

Mehrere GPUs:
Speichern

state = {
        'epoch': epoch,
        'state_dict': model.module.state_dict(),
        'optimizer': optimizer.state_dict(),
}
savepath="checkpoint.t7"
torch.save(state,savepath)

Belastung:

checkpoint = torch.load('checkpoint.t7')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']

#Don't call DataParallel before loading the model otherwise you will get an error

model = nn.DataParallel(model) #ignore the line if you want to load on Single GPU

Lokal speichern

Wie Sie Ihr Modell speichern, hängt davon ab, wie Sie in Zukunft darauf zugreifen möchten. Wenn Sie eine neue Instanz von aufrufen können model Klasse, dann müssen Sie nur noch die Gewichte des Modells speichern/laden 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))

Wenn Sie dies aus irgendeinem Grund nicht können (oder die einfachere Syntax bevorzugen), können Sie das gesamte Modell (eigentlich ein Verweis auf die Datei(en), die das Modell definieren, zusammen mit seinem state_dict) mit speichern torch.save():

# Save:
torch.save(old_model, PATH)

# Load:
new_model = torch.load(PATH)

Da dies jedoch ein Verweis auf den Speicherort der Dateien ist, die die Modellklasse definieren, ist dieser Code nicht portierbar, es sei denn, diese Dateien werden ebenfalls in dieselbe Verzeichnisstruktur portiert.

Speichern in der Cloud – TorchHub

Wenn Sie möchten, dass Ihr Modell portabel ist, können Sie es einfach mit importieren torch.hub. Wenn Sie eine entsprechend definierte hinzufügen hubconf.py Datei in ein Github-Repo, kann dies einfach aus PyTorch heraus aufgerufen werden, damit Benutzer Ihr Modell mit/ohne Gewichte laden können:

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)

Modell laden:

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

pip installiere pytorch-lightning

Stellen Sie sicher, dass Ihr übergeordnetes Modell pl.LightningModule anstelle von nn.Module verwendet

Speichern und Laden von Kontrollpunkten mit Pytorch-Blitz

import pytorch_lightning as pl

model = MyLightningModule(hparams)
trainer.fit(model)
trainer.save_checkpoint("example.ckpt")
new_model = MyModel.load_from_checkpoint(checkpoint_path="example.ckpt")