pytorch基础-保存和加载模型（7）

这一节的内容，将研究如何通过保存、加载和运行预测模型来保持模型状态。导入相应的包：

import torch
import torchvision.models as models

1、 保存和加载模型权重

PyTorch 模型将学习到的参数存储在称为 state_dict 的内部状态字典中。 这些参数可以通过 torch.save 方法保存起来：

model = models.vgg16(pretrained=True)#创建预训练模型vgg16，实例化
torch.save(model.state_dict(), 'model_weights.pth')

 在加载模型权重时，我们需要先实例化模型类，因为该类定义网络的结构。然后使用 load_state_dict() 方法加载参数。

model = models.vgg16() # 没有指定预训练=True，即不加载默认权重
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()#模型中有BatchNormalization和Dropout，在预测时使用model.eval()后会将其关闭以免影响预测结果。

 这里为了方便使用了迁移学习的方法，下面是关于迁移学习方法的具体解释：

一文读懂迁移学习和预训练__Daibingh_的博客-CSDN博客_预训练

 其中torch.save()有两种保存方式：第一种只保存神经网络的训练模型的参数，save的对象是model.state_dict()，第二种既保存整个神经网络的的模型结构又保存模型参数，那么save的对象就是整个模型；

有时可能希望保存此类的结构与模型，在这种情况下可以通过 model（不是 model.state_dict())保存功能： 假设我有一个训练好的模型名叫model，如保存参数以及结构的方法是：

torch.save(model, 'net.pth')  # 保存整个神经网络的模型结构以及参数
torch.save(model.state_dict(), 'net_params.pth')  # 只保存模型参数
model = torch.load('net.pth')  # 加载整个神经网络的模型结构以及参数
model.load_state_dict(torch.load('net_params.pth')) # 仅加载参数

 其中，一定要在测试之前调用 model.eval() 方法，将 dropout 和 batch normalization 层设置为评估模式。 不这样做会产生不一致的推理结果。

如果使用model.eval()，则BN层就不会再计算预测数据的均值和方差，即在预测过程中BN层的均值和方差就是训练过程得到的均值和方差mean_train，variance_train，此时预测结果就不会再发生变化。预测过程中如果不使用model.eval()的话，依然会使一部分的网络连接不进行计算，而使用model.eval()后就是所有的网络连接均进行计算。

2、预加载模型的修改

这里有详细的介绍：

二、现有网络模型（torchvision.models.vgg16）的修改与使用_weixin_44987829的博客-CSDN博客_torchvision vgg16

 （1）修改vgg16_true网路结构，添加linear层：

import torchvision
from torch.nn import Linear

vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True)
print(vgg16_true)
# vgg16_true.add_module("add_linear",Linear(1000,10))
vgg16_true.classifier.add_module("add_linear",Linear(1000,10))
print(vgg16_true)

 输出结果：

  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
    (add_linear): Linear(in_features=1000, out_features=10, bias=True)
  )

 （2）修改vgg16_false网路结构，更改分类器第6层为指定linear层

print(vgg16_false)
vgg16_false.classifier[6]=Linear(4096,10)
print(vgg16_false)

   #更改之前
(6): Linear(in_features=4096, out_features=1000, bias=True)
   #更改之后
(6): Linear(in_features=4096, out_features=10, bias=True)

 3、保存加载自定义的模型

import torch
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )
    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork()
torch.save(model, "model1.pth")
model1 = torch.load("model1.pth")
print(model1)

 可以看到左边目录上出现了保存的模型文件：