PyTorch深度学习模型的保存和加载流程详解
发布日期：2021-12-22 02:56 | 文章来源：脚本之家
一、模型参数的保存和加载
torch.save(module.state_dict(), path)：使用module.state_dict()函数获取各层已经训练好的参数和缓冲区，然后将参数和缓冲区保存到path所指定的文件存放路径（常用文件格式为.pt、.pth或.pkl）。torch.nn.Module.load_state_dict(state_dict)：从state_dict中加载参数和缓冲区到Module及其子类中 。torch.nn.Module.state_dict()函数返回python中的一个OrderedDict类型字典对象，该对象将每一层与它的对应参数和缓冲区建立映射关系，字典的键值是参数或缓冲区的名称。只有那些参数可以训练的层才会被保存到OrderedDict中，例如：卷积层、线性层等。Python中的字典类以“键:值”方式存取数据，OrderedDict是它的一个子类，实现了对字典对象中元素的排序（OrderedDict根据放入元素的先后顺序进行排序）。由于进行了排序，所以顺序不同的两个OrderedDict字典对象会被当做是两个不同的对象。- 示例：
import torch import torch.nn as nn class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 2, 3) self.pool1 = nn.MaxPool2d(2, 2) def forward(self, x): x = self.conv1(x) x = self.pool1(x) return x # 初始化网络 net = Net() net.conv1.weight[0].detach().fill_(1) net.conv1.weight[1].detach().fill_(2) net.conv1.bias.data.detach().zero_() # 获取state_dict state_dict = net.state_dict() # 字典的遍历默认是遍历key，所以param_tensor实际上是键值 for param_tensor in state_dict: print(param_tensor,'：\n',state_dict[param_tensor]) # 保存模型参数 torch.save(state_dict,"net_params.pth") # 通过加载state_dict获取模型参数 net.load_state_dict(state_dict)
输出：
二、完整模型的保存和加载
torch.save(module, path)：将训练完的整个网络模型module保存到path所指定的文件存放路径（常用文件格式为.pt或.pth）。torch.load(path)：加载保存到path中的整个神经网络模型。- 示例：
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 2, 3)
self.pool1 = nn.MaxPool2d(2, 2)
def forward(self, x):
x = self.conv1(x)
x = self.pool1(x)
return x
# 初始化网络
net = Net()
net.conv1.weight[0].detach().fill_(1)
net.conv1.weight[1].detach().fill_(2)
net.conv1.bias.data.detach().zero_()
# 保存整个网络
torch.save(net,"net.pth")
# 加载网络
net = torch.load("net.pth")
到此这篇关于PyTorch深度学习模型的保存和加载流程详解的文章就介绍到这了,更多相关PyTorch 模型的保存 内容请搜索本站以前的文章或继续浏览下面的相关文章希望大家以后多多支持本站！
版权声明：本站文章来源标注为YINGSOO的内容版权均为本站所有，欢迎引用、转载，请保持原文完整并注明来源及原文链接。禁止复制或仿造本网站，禁止在非www.yingsoo.com所属的服务器上建立镜像，否则将依法追究法律责任。本站部分内容来源于网友推荐、互联网收集整理而来，仅供学习参考，不代表本站立场，如有内容涉嫌侵权，请联系alex-e#qq.com处理。
相关文章
关注官方微信