Pytorch學習（五）在一個文件中保存和加載多個模型

保存和加載多個模型有助於重用您之前訓練過的模型。

介紹

在保存由多個torch.nn.Modules組成的模型時。如GAN、sequence-sequence模型或模型集合，必須保存每個模型的state_dict和相應的優化器的字典。您還可以保存其他可能幫助您恢復培訓的項目，只需將它們附加到字典中即可。要加載模型，首先要初始化模型和優化器，然後使用torch.load()本地加載字典。在這裏，只需像預期的那樣查詢字典，就可以輕鬆地訪問保存的項。在這個食譜中，我們將演示如何使用PyTorch將多個模型保存到一個文件中。

步驟

1. 導入包

2. 定義和初始化神經網絡

3. 初始化優化器

4. 存儲多個模型

5. 加載多個模型

1. Import necessary libraries for loading our data

import torch
import torch.nn as nn
import torch.optim as optim

2. Define and intialize the neural network

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

netA = Net()
netB = Net()

3. Initialize the optimizer

optimizerA = optim.SGD(netA.parameters(), lr=0.001, momentum=0.9)
optimizerB = optim.SGD(netB.parameters(), lr=0.001, momentum=0.9)

4. Save multiple models

# Specify a path to save to
PATH = "model.pt"

torch.save({
            'modelA_state_dict': netA.state_dict(),
            'modelB_state_dict': netB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            }, PATH)

5. Load multiple models

記住首先初始化模型和優化器，然後本地加載字典。

modelA = Net()
modelB = Net()
optimModelA = optim.SGD(modelA.parameters(), lr=0.001, momentum=0.9)
optimModelB = optim.SGD(modelB.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

在運行推理之前，必須調用model.eval()將dropout和bath normalization層設置爲計算模式。如果不這樣做，將會產生不一致的推斷結果。如果您希望恢復培訓，請調用model.train()以確保這些層處於培訓模式。恭喜你!您已經成功地在PyTorch中保存並加載了多個模型