pytorch模型导入、修改、保存、读取

pytorch对模型的操作

导入

这里以vgg11的classifier层为例
导入模型:

vgg11 = torchvision.models.vgg11(pretrained=False)
print(vgg11)

其中classifier的输出如下
在这里插入图片描述

修改

主要的修改方式:

# 在某层中添加层
vgg11.classifier.add_module('new_linear', nn.Linear(1000, 10))
# 修改某层
vgg11.classifier[6] = nn.Linear(4096, 10)
print(vgg11)

修改后:
在这里插入图片描述

保存与读取

# 模型的保存与读取
# 方式1,保存模型和参数
torch.save(vgg11, 'vgg11_method1.pth')
# 在读取时需要保证原模型已经引入
model = torch.load('vgg11_method1.pth')

# 方式2,只保存模型参数,一个字典形式(官方推荐)
torch.save(vgg11.state_dict(), 'vgg11_method2.pth')

vgg11_new = torchvision.models.vgg11(pretrained=False)
vgg11_new.load_state_dict(torch.load('vgg11_method2.pth'))
  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值