第一、实现模型保存的意义

大家应该都遇到过，在模型训练过程中会发生意外中断(例如断电)的情况。假设我们需要训练10000个epoch(迭代训练)，当我们训练到第9990次时突然中断了，而且整个模型也没有得到保存，则我们的模型需要重新开始训练。众所周知，训练一个模型需要耗费大量的电能和GPU资源。因此在训练完成之前，我们需要设定在每一个或者多个epoch训练结束后，保存一次模型。一方面可以防止断电重跑，另一方面可以观察不同迭代次数模型的表现；在训练完成以后，我们需要保存模型参数值用于后续的测试过程。所以，保存的对象包含网络参数值、优化器参数值、epoch值等等。

第二、定义一个简单的网络

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

这是一个简单的卷积#神经网络#，其中Conv是卷积层，Pool是池化层，fc是全连接层，其中的参数、代码就不详细讲了。该网络的功能是实现10个类别(如下图代码)的分类，简单来说是输入一张3通道(#RGB#)图像，输出10个值，值最大的是图像中的类别。

classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')

第三、训练模型

    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=10, shuffle=True, num_workers=5)
    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=10, shuffle=False, num_workers=5)
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    #上面是定义训练集和测试集
    #下面开始定义模型、损失函数
    net = Net()
    optimizer=SGD(net.parameters(), lr=0.001, momentum=0.9)  #随机梯度下降法
    epoches=0
   
    for epoch in range(epoches+1,10000):  # epoch从1到10000
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):  #枚举函数用来生成迭代器，迭代对象trainloader和起始位置start=0 
            inputs, labels = data    #输入图片
            optimizer.zero_grad()   #梯度设为0
            outputs = net(inputs) #模型的输出值
            loss = criterion(outputs, labels) #通过损失函数来计算模型输出值和实际值之间的损失值
            loss.backward()       #损失值反向传播
            optimizer.step()      #更新所有的参数
            #下面是输出当前状态(epoch值、损失值)
            print('epoch:%d,  loss: %.3f' %(epoch + 1, loss.item()))
    print('Finished Training')  #提示训练完成

上面展示了如何简单训练一个模型的代码过程，我们可以清楚地看到，在10000个epoch训练完成前，无论出现任何意外导致中断的情况，都需要重新开始训练，所以保存模型就变得尤为重要。

第四、保存模型的代码块

1.保存Net的参数值

其中state_dict 是一个简单的python的字典对象,将每一层与它的对应参数建立映射关系(如net的每一层的w权重和偏置；只有那些参数可以训练的层才会被保存到模型中，例如卷积层,池化层和全连接层等；优化器Optimizer也有一个state_dict,它包含了优化器的状态、被使用的超参数，如lr, momentum,weight_decay等)。

f后面应该跟着一个保存文件的路径，假设保存到同级目录models下，并命名为net.pth，运行代码后，在./models/将会出现net.pth文件，pth文件中的内容就是net的参数名称和值对应的字典。

2.加载Net的参数值

net = Net()
net.load_state_dict(torch.load("models/net.pth"))

通过上面语句就能实现上面保存模型参数值的加载。

3.优化器与epoch值的保存

上面保存模型的简单参数通过两三行#代码#即可完成，现在加上优化器与epoch值，需要创建一个字典来保存所有的对象，同样也用save函数保存这个字典。(下图代码也是保存上面代码模型的参数)

all_states = {"net": net.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch}
torch.save(obj=all_states, f="models/net.pth")

第一条语句创建了一个字典(all_states)来保存网络参数值、优化器和epoch值；第二条语句和上面的save方式一致。为了减少损失，建议每隔一段时间保存一次参数。

4.优化器与epoch值的加载

加载语句会比上面的多一些，因为要加载后要对3个参数赋值。

    net = Net()
    optimizer=SGD(net.parameters(), lr=0, momentum=0)
    all_states=torch.load("./models/net.pth")
    optimizer.load_state_dict(all_states["optimizer"])
    epoches=all_states["epoch"]
    net.load_state_dict(all_states["net"])

第五、总结

这些步骤应该都是通俗易懂的，当中最重要的部分是“字典”的概念和操作，因为参数的存储是需要“名称”和“值”对应，读取时也是通过键值对读取的。