保存和加载pytorch模型

当保存和加载模型时，需要熟悉三个核心功能：

torch.save：将序列化对象保存到磁盘。此函数使用Python的pickle模块进行序列化。使用此函数可以保存如模型、tensor、字典等各种对象。
torch.load：使用pickle的unpickling功能将pickle对象文件反序列化到内存。此功能还可以有助于设备加载数据。
torch.nn.Module.load_state_dict：使用反序列化函数 state_dict 来加载模型的参数字典。

Python中对于模型数据的保存和加载操作都是引用Python内置的pickle包，使用pickle.dump()和pickle.load()方法。在Pytorch中也有同样功能的方法提供。

>>>torch.save(model,'model.pkl') #保存整个模型
>>>model = torch.load('model.pkl') #加载整个模型
>>>torch.save(alexnet.state_dict(),'params.pkl') #保存网格中的参数
>>>alexnet.load_state_dict(torch.load('params.pkl')) #加载网格中的参数

在torchvision.models模块里，PyTorch提供了一些常用的模型：

常用模型
AlexNet
VGG
ResNet
SqueezeNet
DenseNet
Inception v3

可以使用torch.util.model_zoo来预加载它们，具体设置通过参数pretrained=True来实现。

>>>import torchvision.models as models
>>>ResNet18 = models.ResNet18(pretrained=True)
>>>alexnet = models.alexnet(pretrained=True)
>>>squeezenet = models.squeezenet1_0(pretrained=True)
>>>vgg16 = models.vgg16(pretrained=True)
>>>densenet = models.densenet161(pretrained=True)
>>>inception = models.inception_v3(pretrained=True)

加载这类预训练模型的过程中，还可以进行微处理。

>>>pretrained_dict = model_zoo.load_url(model_urls['resnet134'])
>>>model_dict = model.state_dict()
>>>pretrained_dict = {k:v for k,v in pretrained_dict.items()if k in model_dict}#将pretrained_dict里不属于model_dict的键剔除掉
>>>model_dict.update(pretrained_dict)  #更新现在有的model_dict
>>>model.load_state_dict(model_dict)

参考
《PyTorch机器学习从入门到实战》

保存和加载pytorch模型相关推荐

加载dict_PyTorch 7.保存和加载pytorch模型的两种方法
众所周知,python的对象都可以通过torch.save和torch.load函数进行保存和加载(不知道?那你现在知道了(*^_^*)),比如: x1 = {"d":" ...
PyTorch 深度剖析：如何保存和加载PyTorch模型？
点击上方"视学算法",选择加"星标"或"置顶" 重磅干货,第一时间送达作者丨科技猛兽编辑丨极市平台导读本文详解了PyTorch 模型 ...
pytorch模型的保存和加载、checkpoint
pytorch模型的保存和加载.checkpoint 其实之前笔者写代码的时候用到模型的保存和加载,需要用的时候就去度娘搜一下大致代码,现在有时间就来整理下整个pytorch模型的保存和加载,开始学习 ...
TensorFlow模型保存和加载方法
TensorFlow模型保存和加载方法模型保存 import tensorflow as tfw1 = tf.Variable(tf.constant(2.0, shape=[1]), name=& ...
keras模型保存和加载
(一)保存和加载整个模型包含模型的结构.权重.训练配置项(损失函数.优化器).优化器状态,允许准确地从上次结束的地方开始训练. 1.训练完模型后 path='.../.../xxx.h5' mo ...
PyTorch | 保存和加载模型教程
点击上方"算法猿的成长",选择"加为星标" 第一时间关注 AI 和 Python 知识图片来自 Unsplash,作者: Jenny Caywood 2019 ...
python保存模型与参数_基于pytorch的保存和加载模型参数的方法
当我们花费大量的精力训练完网络,下次预测数据时不想再(有时也不必再)训练一次时,这时候torch.save(),torch.load()就要登场了. 保存和加载模型参数有两种方式: 方式一: torc ...
PyTorch学习笔记2：nn.Module、优化器、模型的保存和加载、TensorBoard
文章目录一.nn.Module 1.1 nn.Module的调用 1.2 线性回归的实现二.损失函数三.优化器 3.1.1 SGD优化器 3.1.2 Adagrad优化器 3.2 分层学习率 3 ...
pytorch load state dict_Pytorch学习记录-使用Pytorch进行深度学习，保存和加载模型
新建 Microsoft PowerPoint 演示文稿 (2).jpg 保存和加载模型在完成60分钟入门之后,接下来有六节tutorials和五节关于文本处理的tutorials.争取一天一节.不 ...

保存和加载pytorch模型

保存和加载pytorch模型相关推荐

最新文章

热门文章