模型的保存与加载(pytorch)

(联邦学习笔记,资料来源于b站小土堆)

对于我们自定义的或者对于现存的模型进行修改后,我们可能要在别的地方进行调用,这时候需要先将它进行保存以便后续调用。

模型保存

对于模型的保存有两种方式,一种是同时保存模型结构和模型参数;另一种是仅保存模型参数(官方推荐)

保存方式一(保存模型结构和模型参数):

torch.save(vgg16,"vgg16_methond1.pth")

其中vgg16是模型,vgg16_methond1.pth是保存的名称

保存方式二(仅保存模型参数):

torch.save(vgg16.state_dict(),"vgg16_methond2.pth")

如果要保存自定义的模型,如下操作:

#保存自定义的模型
class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.conv1 = nn.Conv2d(3,64,kernel_size=3)

    def forward(self,input):
        output = self.conv1(input)
        return output

mymodule = MyModule()

torch.save(mymodule,"mymodule_save.pth")

以上是保存的相关操作,可定义在一个python文件中。完整代码如下:

import torch
import torchvision
from torch import nn

vgg16 = torchvision.models.vgg16(pretrained=False)


#保存方式一,保存模型结构和模型参数
torch.save(vgg16,"vgg16_methond1.pth")

#保存方式二,仅保存模型参数(官方推荐)
torch.save(vgg16.state_dict(),"vgg16_methond2.pth")

#保存自定义的模型
class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.conv1 = nn.Conv2d(3,64,kernel_size=3)

    def forward(self,input):
        output = self.conv1(input)
        return output

mymodule = MyModule()

torch.save(mymodule,"mymodule_save.pth")

模型加载

加载内容一般根据的保存的方式来显示,一般形式为:model = torch.load("vgg16_methond1.pth")

对于上面的保存方式二,如果想要以(模型结构+模型参数)方式加载,则要先知道原本的模型是什么(上述是vgg16),然后进行操作,具体如下:

#加载方式2 同时加载模型和参数
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_methond2.pth"))
print(vgg16)

模型加载可写在同一个python文件中,完整代码如下:

import torch
import torchvision

import model_save


#加载以 方式一 保存的模型(模型结构和参数)
model = torch.load("vgg16_methond1.pth")
print(model)

#加载 方式二 保存的模型(模型参数)

#加载方式1 直接加载模型参数
model1 = torch.load("vgg16_methond2.pth")
print(model1)

#加载方式2 同时加载模型和参数
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_methond2.pth"))
print(vgg16)

#加载自定义模型
model2 = torch.load("mymodule_save.pth")
print(model2)

(如有不同意见,欢迎留下评论,我见到了会第一时间回复)

  • 1
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值