Pytorch模型参数的保存和加载
一、前言
在模型训练完成后，我们需要保存模型参数值用于后续的测试过程。由于保存整个模型将耗费大量的存储，故推荐的做法是只保存参数，使用时只需在建好模型的基础上加载。
通常来说，保存的对象包括网络参数值、优化器参数值、epoch值等。本文将简单介绍保存和加载模型参数的方法，同时也给出保存整个模型的方法供大家参考。
二、参数保存
在这里我们使用 torch.save() 函数保存模型参数：
import torch path = './model.pth' torch.save(model.state_dict(), path)
model——指定义的模型实例变量，如model=net( )
state_dict()——state_dict( )是一个可以轻松地保存、更新、修改和恢复的python字典对象, 对于model来说，表示模型的每一层的权重及偏置等参数信息；对于 optimizer 来说，其包含了优化器的状态以及被使用的超参数(如lr, momentum,weight_decay等)
path——path是保存参数的路径，一般设置为 path='./model.pth' , path='./model.pkl'等形式。
此外，如果想保存某一次训练采用的optimizer、epochs等信息，可将这些信息组合起来构成一个字典保存起来：
import torch
path = './model.pth'
state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
torch.save(state, path)
三、参数的加载
使用 load_state_dict()函数加载参数到模型中， 当仅保存了模型参数，而没有optimizer、epochs等信息时：
model.load_state_dict(torch.load(path))
model——事先定义好的跟原模型一致的模型
path——之前保存的模型参数文件
如若保存了optimizer、epochs等信息，我们这样载入信息：
# 使用torch.load()函数将文件中字典信息载入 state_dict 变量中 state_dict = torch.load(path) # 分布加载参数到模型和优化器 model.load_state_dict(state_dict['model']) optimizer.load_state_dict(state_dict['optimizer']) epoch = state_dict(['epoch'])
我们还可以在每n个epoch后保存一次参数，以观察不同迭代次数模型的表现。此时我们可设置不同的path，如 path='./model' + str(epoch) +'.pth'，这样，不同epoch的参数就能保存在不同的文件中。
四、保存和加载整个模型
使用上文提到的方法即可：
torch.save(model, path) model = torch.load(path)
五、总结
pytorch中state_dict()和load_state_dict()函数配合使用可以实现状态的获取与重载，load()和save()函数配合使用可以实现参数的存储与读取。掌握对应的函数使用方法就可以游刃有余地进行运用。
到此这篇关于Pytorch模型参数的保存和加载的文章就介绍到这了,更多相关Pytorch模型参数保存内容请搜索猪先飞以前的文章或继续浏览下面的相关文章希望大家以后多多支持猪先飞！
原文出处：https://blog.csdn.net/weixin_45684362/article/details/129446
相关文章
pytorch nn.Conv2d()中的padding以及输出大小方式
今天小编就为大家分享一篇pytorch nn.Conv2d()中的padding以及输出大小方式，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-04-27- 这篇文章主要介绍了PyTorch一小时掌握之迁移学习篇，本文给大家介绍的非常详细，对大家的学习或工作具有一定的参考借鉴价值，需要的朋友可以参考下...2021-09-08
Linux安装Pytorch1.8GPU(CUDA11.1)的实现
这篇文章主要介绍了Linux安装Pytorch1.8GPU(CUDA11.1)的实现，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧...2021-03-25- 这篇文章主要介绍了Pytorch之扩充tensor的操作，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2021-03-05
- 今天小编就为大家分享一篇pytorch 自定义卷积核进行卷积操作方式，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-05-06
- 这篇文章主要介绍了解决pytorch 交叉熵损失输出为负数的问题，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-07-08
pytorch绘制并显示loss曲线和acc曲线，LeNet5识别图像准确率
今天小编就为大家分享一篇pytorch绘制并显示loss曲线和acc曲线，LeNet5识别图像准确率，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-05-02- 这篇文章主要介绍了pytorch 实现冻结部分参数训练另一部分，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2021-03-27
从Pytorch模型pth文件中读取参数成numpy矩阵的操作
这篇文章主要介绍了从Pytorch模型pth文件中读取参数成numpy矩阵的操作，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2021-03-04Pytorch 的损失函数Loss function使用详解
今天小编就为大家分享一篇Pytorch 的损失函数Loss function使用详解，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-05-02- 今天小编就为大家分享一篇pytorch中的上采样以及各种反操作,求逆操作详解，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-04-30
- 这篇文章主要介绍了基于Pytorch版yolov5的滑块验证码破解思路详解,本文给大家介绍的非常详细，对大家的学习或工作具有一定的参考借鉴价值，需要的朋友可以参考下...2021-02-25
pytorch中交叉熵损失(nn.CrossEntropyLoss())的计算过程详解
今天小编就为大家分享一篇pytorch中交叉熵损失(nn.CrossEntropyLoss())的计算过程详解，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-05-02- 这篇文章主要介绍了pytorch深度学习中对softmax实现进行了详细解析，有需要的朋友可以借鉴参考下，希望能够有所帮助，祝大家多多进步...2021-09-30
- 今天小编就为大家分享一篇Pytorch 计算误判率,计算准确率,计算召回率的例子，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-04-27
- 今天小编就为大家分享一篇Pytorch实现LSTM和GRU示例，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-04-27
- 这篇文章主要介绍了Pytorch如何切换 cpu和gpu的使用详解，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧...2021-03-01
- 今天小编就为大家分享一篇pytorch动态网络以及权重共享实例，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧...2020-04-29
- 这篇文章主要介绍了解决Pytorch修改预训练模型时遇到key不匹配的情况，具有很好的参考价值，希望对大家有所帮助。如有错误或未考虑完全的地方，望不吝赐教...2021-06-05
- 这篇文章主要介绍了pytorch中的squeeze函数、cat函数使用，具有很好的参考价值，希望对大家有所帮助。如有错误或未考虑完全的地方，望不吝赐教...2021-05-20