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中保存并加载了多个模型