模型的保存与加载
首先，需要导入两个包
import torch import torchvision.models as models
保存和加载模型参数
PyTorch模型将学习到的参数存储在一个内部状态字典中，叫做state_dict。这可以通过torch.save方法来实现。
我们导入预训练好的VGG16模型，并将其保存。我们将state_dict字典保存在model_weights.pth文件中。
model = models.vgg16(pretrained=True) torch.save(model.state_dict(), 'model_weights.pth')
想要加载模型参数，我们需要创建一个和原模型一样的实例，然后通过load_state_dict()方法来加载模型参数
- 创建一个
VGG16模型实例（未经过预训练的） - 加载本地参数
model = models.vgg16() # we do not specify pretrained=True, i.e. do not load default weights
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()
注意：在进行测试前，如果模型中有dropout层和batch normalization层的话，一定要使用model.eval()将模型转到测试模式。
- 在
train模式下，dropout网络层会按照设定的参数p设置保留激活单元的概率（保留概率=p);batchnorm层会继续计算数据的mean和var等参数并更新。 - 在
val模式下，dropout层会让所有的激活单元都通过，而batchnorm层会停止计算和更新mean和var，直接使用在训练阶段已经学出的mean和var值
当然，相同的，在模型进行训练之前，要使用model.train()来将模型转为训练模式
保存和加载模型参数与结构
当加载模型权重时，我们需要首先实例化模型类，因为类定义了网络的结构。我们可能希望将这个类的结构与模型保存在一起。这样的话，我们可以将model而不是model.state_dict()作为参数。
torch.save(model, 'model.pth')
这样，我们加载模型的时候就不用再新建一个实例了。加载方式如下所示
model = torch.load('model.pth')这种方式在网络比较大的时候可能比较慢，因为相较于上面的方式多存储了网络的结构
总结
到此这篇关于PyTorch模型的保存与加载方法的文章就介绍到这了,更多相关PyTorch模型保存加载内容请搜索Devmax以前的文章或继续浏览下面的相关文章希望大家以后多多支持Devmax！