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

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 nnclass 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 nnclass 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模型的保存和加载、checkpoint
pytorch模型的保存和加载.checkpoint 其实之前笔者写代码的时候用到模型的保存和加载,需要用的时候就去度娘搜一下大致代码,现在有时间就来整理下整个pytorch模型的保存和加载,开始学习 ...
PyTorch模型的保存加载以及数据的可视化
文章目录 PyTorch模型的保存和加载模块和张量的序列化和反序列化模块状态字典的保存和载入 PyTorch数据的可视化 TensorBoard的使用总结 PyTorch模型的保存和加载在深度 ...
PyTorch学习笔记2：nn.Module、优化器、模型的保存和加载、TensorBoard
文章目录一.nn.Module 1.1 nn.Module的调用 1.2 线性回归的实现二.损失函数三.优化器 3.1.1 SGD优化器 3.1.2 Adagrad优化器 3.2 分层学习率 3 ...
numpy将所有数据变为0和1_PyTorch 学习笔记（二）：张量、变量、数据集的读取、模组、优化、模型的保存和加载...
一. 张量 PyTorch里面最基本的操作对象就是Tensor,Tensor是张量的英文,表示的是一个多维的矩阵,比如零维就是一个点,一维就是向量,二维就是一般的矩阵,多维就相当于一个多维的数组,这和 ...
线性回归之模型的保存和加载
线性回归之模型的保存和加载 1 sklearn模型的保存和加载API from sklearn.externals import joblib [目前这行代码报错,直接写import joblib ...
paddlepaddle模型的保存和加载
导读深度学习中模型的计算图可以被分为两种,静态图和动态图,这两种模型的计算图各有优劣. 静态图需要我们先定义好网络的结构,然后再进行计算,所以静态图的计算速度快,但是debug比较的困难,因为只有当 ...
PyTorch基础-模型的保存和加载-09
模型的保存 import numpy as np import torch from torch import nn,optim from torch.autograd import Variable ...
tensorflow 模型的保存和加载
为了让训练结果可以复用,需要将训练得到的神经网络模型持久化,也就是把模型的参数保存下来,并保证可以持久化后的模型文件中还原出保存的模型. 1. 保存模型 tensorflow提供了一个API可以方便的 ...
Pytorch模型训练保存/加载（搭建完整流程）
文章目录前言模型训练完整步骤模型保存与加载 GPU训练 "借鸡生蛋" 模型使用本博文优先在掘金社区发布! 前言我们这边还是以CIARF10这个模型为例子. 现在的话先说明 ...

PyTorch | 模型的保存和加载

PyTorch | 模型的保存和加载

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

二、完整模型的保存和加载

PyTorch | 模型的保存和加载相关推荐

最新文章

热门文章