Pytorch保存和加载模型(load和load_state_dict)

本文介绍了PyTorch中模型的保存和加载,重点讲解了state_dict的概念和使用。通过对比两种保存模型的方法,阐述了当只保存state_dict时,需要先构建模型结构再加载参数。实验展示了只保存权重(only_weights.pth)和保存完整模型(whole_model.pth)的区别,强调了在没有模型结构的情况下尝试加载权重会导致错误。
摘要由CSDN通过智能技术生成

Pytorch目前成为学术界最流行的DL框架,没有之一。很大程度上,简洁直观地操作有关。模型的保存和加载,于pytorch而言,也是很简单的。本文做了一个比较实验,方便大家理解。

首先,要清楚几个函数:torch.savetorch.loadstate_dict()load_state_dict()
先举最简单的例子:

import torch

model = torch.load('my_model.pth')
torch.save(model, 'new_model.pth')

上面的代码非常直观,一载一存。但是有一个问题,这样保存的pth文件直接包含了整个模型的结构。当你需要灵活加载模型参数时,比如只加载部分参数,那么这种情况保存的pth文件读取进来还得额外解析出“参数文件”


如果想更灵活对待咱们训练好的模型参数,咱们可以使用下面这个方法。pytorch把所有的模型参数用一个内部定义的dict进行保存,自称为“state_dict。这个所谓的state_dict就是不带模型结构的模型参数了~
咱们的加载和保存就发生了一点微妙的变化:

import torch
model = MyModel() # init your model class, build the graph shape

state_dict = torch.load('model_state_dict.pth')
model.load_state_dict(state_dict)

torch.save(model.state_dict(), 'model_state_dict1.pth')

比较上面两段代码,咱们可以有一下结论:

  1. pth文件既可能保存了模型的图结构,也有可能没保存;
  2. 加载没保存图结构的pth时,需要先初始化模型结构,即把架子搭好;
  3. 在保存模型的时候,如果不想保存图结构,可以单独保存model.state_dict()

到这里还不清楚的可以留言提问~
参考链接:https://pytorch.org/tutorials/beginner/basics/saveloadrun_tutorial.html


实验

脚本如下:

import torch
import torchvision.models as models

model = models.vgg16(pretrained=True)
torch.save(model.state_dict(), 'only_weights.pth')

model_state_dict = torch.load('only_weights.pth')
model1 = models.vgg16() # describe the graph shape
model1.load_state_dict(model_state_dict)
model1.eval()

torch.save(model1, 'whole_model.pth')

model2 = torch.load('whole_model.pth')
model2.eval()

# model3 = torch.load('only_weights.pth')
# model3.eval()    # Error

model3切换到eval()模式就会报错,原因是model3只包含weights而缺乏图结构~

评论 6
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

木盏

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值