Python 官方文档:入门教程 => 点击学习
目录模型的保存与加载保存和加载模型参数保存和加载模型参数与结构总结模型的保存与加载 首先,需要导入两个包 import torch import torchvision.models
首先,需要导入两个包
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模型保存加载内容请搜索编程网以前的文章或继续浏览下面的相关文章希望大家以后多多支持编程网!
--结束END--
本文标题: PyTorch模型的保存与加载方法实例
本文链接: https://www.lsjlt.com/news/120553.html(转载时请注明来源链接)
有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341
下载Word文档到电脑,方便收藏和打印~
2024-03-01
2024-03-01
2024-03-01
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
回答
回答
回答
回答
回答
回答
回答
回答
回答
回答
0