Speichern und lesen Sie das Netzwerkmodell von Torch

Speichern und lesen Sie das Netzwerkmodell von Torch

  Heute verwenden wir zwei Methoden, um das Speichern und Lesen des Modells zu veranschaulichen: Es wird bestätigt, dass die Speichermethode des Modells mit der Lesemethode übereinstimmen muss. Beispiel: Das Modell, das wir in Methode 1 gespeichert haben, muss beim Lesen in Methode 1 gelesen werden.

1. Speichern Sie das Modell

vgg16 = torchvision.models.vgg16(pretrained=False)
#方式一：
torch.save(vgg16, "vgg16_method1.pth")
#方式二：
torch.save(vgg16.state_dict(), "vgg16_method2.pth")

2. Modelllesung

#方式一：
model = torch.load("vgg16_method1.pth")
#方式二：
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))