PyTorch模型的保存与加载方法实例

更新时间：2022年09月23日 09:53:59 作者：长命百岁️

Pytorch保存模型其实非常简单,下面这篇文章主要给大家介绍了关于PyTorch模型的保存与加载的相关资料,文中通过实例代码介绍的非常详细,需要的朋友可以参考下

模型的保存与加载

首先，需要导入两个包

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模型保存加载内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！

您可能感兴趣的文章:

实例讲解Python 迭代器与生成器
这篇文章主要介绍了Python 迭代器与生成器的相关资料，文中示例代码非常详细，帮助大家更好的理解和学习，感兴趣的朋友可以了解下
2020-07-07
Python中GPU计算的库pycuda的使用
本文主要介绍了Python中GPU计算的库pycuda的使用,详细介绍了PyCUDA 库的特性、用法,并通过丰富的示例代码展示其在实际项目中的应用,感兴趣的可以了解一下
2024-05-05
对DataFrame数据中的重复行,利用groupby累加合并的方法详解
今天小编就为大家分享一篇对DataFrame数据中的重复行,利用groupby累加合并的方法详解，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2019-01-01
解析ROC曲线绘制(python+sklearn+多分类)
这篇文章主要介绍了解析ROC曲线绘制(python+sklearn+多分类)，具有很好的参考价值，希望对大家有所帮助。如有错误或未考虑完全的地方，望不吝赐教
2022-11-11
Python随机生成一个6位的验证码代码分享
这篇文章主要介绍了Python随机生成一个6位的验证码代码分享,本文直接给出代码实例,需要的朋友可以参考下
2015-03-03
使用Matplotlib绘制平行坐标系的示例详解
平行坐标系，是一种含有多个垂直平行坐标轴的统计图表，这篇文章主要为大家介绍了如何使用绘制平行坐标系，需要的小伙伴可以参考一下
2023-07-07
Python中enumerate()函数详细分析(附多个Demo)
Python的enumerate()函数是一个内置函数,主要用于在遍历循环中获取每个元素的索引以及对应的值,这篇文章主要介绍了Python中enumerate()函数的相关资料,需要的朋友可以参考下
2024-10-10
python设置中文界面实例方法
在本篇文章里小编给大家整理的是关于python设置中文界面实例方法，有兴趣的朋友们可以学习参考下。
2020-10-10
Python功能键的读取方法
这篇文章主要介绍了Python功能键的读取方法,涉及Python键盘事件的相关操作技巧,需要的朋友可以参考下
2015-05-05
Python中卷积神经网络(CNN)入门教程分分享
卷积神经网络（Convolutional Neural Networks, CNN）是一类特别适用于处理图像数据的深度学习模型，本文介绍了如何使用Keras创建一个简单的CNN模型，并用它对手写数字进行分类，需要的可以参考一下
2023-05-05