这一节的内容,将研究如何通过保存、加载和运行预测模型来保持模型状态。导入相应的包:
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()后会将其关闭以免影响预测结果。
这里为了方便使用了迁移学习的方法,下面是关于迁移学习方法的具体解释:
其中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)的修改与使用
(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)
可以看到左边目录上出现了保存的模型文件: