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

到此这篇关于PyTorch深度学习模型的保存和加载流程详解的文章就介绍到这了,更多相关PyTorch 模型的保存 内容请搜索我们以前的文章或继续浏览下面的相关文章希望大家以后多多支持我们！

赞 (0)

相关推荐

  • Pytorch 统计模型参数量的操作 param.numel()

    Pytorch 统计模型参数量的操作 param.numel()

    param.numel() 返回param中元素的数量 统计模型参数量 num_params = sum(param.numel() for param in net.parameters()) print(num_params) 补充:Pytorch 查看模型参数 Pytorch 查看模型参数 查看利用Pytorch搭建模型的参数,直接看程序 import torch # 引入torch.nn并指定别名 import torch.nn as nn import torch.nn.functio

  • pytorch模型的保存和加载、checkpoint操作

    其实之前笔者写代码的时候用到模型的保存和加载,需要用的时候就去度娘搜一下大致代码,现在有时间就来整理下整个pytorch模型的保存和加载,开始学习把~ pytorch的模型和参数是分开的,可以分别保存或加载模型和参数.所以pytorch的保存和加载对应存在两种方式: 1. 直接保存加载模型 (1)保存和加载整个模型 # 保存模型 torch.save(model, 'model.pth\pkl\pt') #一般形式torch.save(net, PATH) # 加载模型 model = torc

  • pytorch 预训练模型读取修改相关参数的填坑问题

    pytorch 预训练模型读取修改相关参数的填坑 修改部分层,仍然调用之前的模型参数. resnet = resnet50(pretrained=False) resnet.load_state_dict(torch.load(args.predir)) res_conv31 = Bottleneck_dilated(1024, 256,dilated_rate = 2) print("---------------------",res_conv31) print("---

  • 关于Pytorch中模型的保存与迁移问题

    目录 1 引言 2 模型的保存与复用 2.1 查看网络模型参数 2.2 载入模型进行推断 2.3 载入模型进行训练 2.4 载入模型进行迁移 3 总结 1 引言 各位朋友大家好,欢迎来到月来客栈.今天要和大家介绍的内容是如何在Pytorch框架中对模型进行保存和载入.以及模型的迁移和再训练.一般来说,最常见的场景就是模型完成训练后的推断过程.一个网络模型在完成训练后通常都需要对新样本进行预测,此时就只需要构建模型的前向传播过程,然后载入已训练好的参数初始化网络即可. 第2个场景就是模型的再训练过

  • PyTorch零基础入门之构建模型基础

    PyTorch零基础入门之构建模型基础

    目录 一.神经网络的构造 二.神经网络中常见的层 2.1 不含模型参数的层 2.2 含模型参数的层 (1)代码栗子1 (2)代码栗子2 2.3 二维卷积层 stride 2.4 池化层 三.LeNet模型栗子 三点提醒: 四.AlexNet模型栗子 Reference 一.神经网络的构造 PyTorch中神经网络构造一般是基于 Module 类的模型来完成的,它让模型构造更加灵活.Module 类是 nn 模块里提供的一个模型构造类,是所有神经网络模块的基类,我们可以继承它来定义我们想要的模型.

  • 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()函数

  • Tensorflow2.1 完成权重或模型的保存和加载

    目录 前言 实现方法 1. 读取数据 2. 搭建深度学习模型 3. 使用回调函数在每个 epoch 后自动保存模型权重 4. 使用回调函数每经过 5 个 epoch 对模型权重保存一次 5. 手动保存模型权重到指定目录 6. 手动保存整个模型结构和权重 前言 本文主要使用 cpu 版本的 tensorflow-2.1 来完成深度学习权重参数/模型的保存和加载操作. 在我们进行项目期间,很多时候都要在模型训练期间.训练结束之后对模型或者模型权重进行保存,然后我们可以从之前停止的地方恢复原模型效果继

  • 解决tensorflow模型参数保存和加载的问题

    终于找到bug原因!记一下:还是不熟悉平台的原因造成的! Q:为什么会出现两个模型对象在同一个文件中一起运行,当直接读取他们分开运行时训练出来的模型会出错,而且总是有一个正确,一个读取错误? 而 直接在同一个文件又训练又重新加载模型预测不出错,而且更诡异的是此时用分文件里的对象加载模型不会出错? model.py,里面含有 ModelV 和 ModelP,另外还有 modelP.py 和 modelV.py 分别只含有 ModelP 和 ModeV 这两个对象,先使用 modelP.py 和 m

  • Python深度学习之图像标签标注软件labelme详解

    Python深度学习之图像标签标注软件labelme详解

    前言 labelme是一个非常好用的免费的标注软件,博主看了很多其他的博客,有的直接是翻译稿,有的不全面.对于新手入门还是有点困难.因此,本文的主要是详细介绍labelme该如何使用. 一.labelme是什么? labelme是图形图像注释工具,它是用Python编写的,并将Qt用于其图形界面.说直白点,它是有界面的, 像软件一样,可以交互,但是它又是由命令行启动的,比软件的使用稍微麻烦点.其界面如下图: 它的功能很多,包括: 对图像进行多边形,矩形,圆形,多段线,线段,点形式的标注(可用于目

  • SpringIOC BeanDefinition的加载流程详解

    SpringIOC BeanDefinition的加载流程详解

    目录 一.前言 二. BeanDefinition 的体系 2.1 体系概览 2.2 BeanDefinition 的作用 三. BeanDefinition 的载入 3.1 载入的入口 3.2 保存的逻辑 3.3 使用的方式 总结 一.前言 这一篇来看看 SpringIOC 里面的一个细节点 , 来简单看看 BeanDefinition 这个对象 , 以及有没有办法对其进行定制. CASE 备份 :  gitee.com/antblack/ca… 二. BeanDefinition 的体系 2

  • Vue openLayers实现图层数据切换与加载流程详解

    Vue openLayers实现图层数据切换与加载流程详解

    目录 openlayers介绍 一.实现效果预览 二.代码实现 openlayers介绍 OpenLayers是一个用于开发WebGIS客户端的JavaScript包.OpenLayers 支持的地图来源包括Google Maps.Yahoo. Map.微软Virtual Earth 等,用户还可以用简单的图片地图作为背景图,与其他的图层在OpenLayers 中进行叠加,在这一方面OpenLayers提供了非常多的选择.OpenLayers采用面向对象方式开发. OpenLayers 是一个专

  • Python深度学习实战PyQt5布局管理项目示例详解

    Python深度学习实战PyQt5布局管理项目示例详解

    目录 1. 从绝对定位到布局管理 1.1 什么是布局管理 1.2 Qt 中的布局管理方法 2. 水平布局(Horizontal Layout) 3. 垂直布局(Vertical Layout) 4. 栅格布局(Grid Layout) 5. 表格布局(Form Layout) 6. 嵌套布局 7. 容器布局 布局管理就是管理图形窗口中各个部件的位置和排列.图形窗口中的大量部件也需要通过布局管理,对部件进行整理分组.排列定位,才能使界面整齐有序.美观大方. 1. 从绝对定位到布局管理 1.1 什么

  • 深度学习小工程练习之tensorflow垃圾分类详解

    深度学习小工程练习之tensorflow垃圾分类详解

    介绍 这是一个基于深度学习的垃圾分类小工程,用深度残差网络构建 软件架构 使用深度残差网络resnet50作为基石,在后续添加需要的层以适应不同的分类任务 模型的训练需要用生成器将数据集循环写入内存,同时图像增强以泛化模型 使用不包含网络输出部分的resnet50权重文件进行迁移学习,只训练我们在5个stage后增加的层 安装教程 需要的第三方库主要有tensorflow1.x,keras,opencv,Pillow,scikit-learn,numpy 安装方式很简单,打开terminal,例

  • 基于pytorch的保存和加载模型参数的方法

    当我们花费大量的精力训练完网络,下次预测数据时不想再(有时也不必再)训练一次时,这时候torch.save(),torch.load()就要登场了. 保存和加载模型参数有两种方式: 方式一: torch.save(net.state_dict(),path): 功能:保存训练完的网络的各层参数(即weights和bias) 其中:net.state_dict()获取各层参数,path是文件存放路径(通常保存文件格式为.pt或.pth) net2.load_state_dict(torch.loa

随机推荐