PyTorch 模型的保存 PyTorch深度学习模型的保存和加载流程详解

软件发布|下载排行|最新软件

当前位置:首页IT学院IT技术

PyTorch 模型的保存 PyTorch深度学习模型的保存和加载流程详解

软耳朵DONG   2021-10-21 我要评论
想了解PyTorch深度学习模型的保存和加载流程详解的相关内容吗,软耳朵DONG在本文为您仔细讲解PyTorch 模型的保存 的相关知识和一些Code实例,欢迎阅读和指正,我们先划重点:PyTorch,模型的保存,PyTorch,模型的加载,下面大家一起来学习吧。

一、模型参数的保存和加载

  •  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")

Copyright 2022 版权所有 软件发布 访问手机版

声明:所有软件和文章来自软件开发商或者作者 如有异议 请与本站联系 联系我们