[PyTorch 学习笔记] 7.1 模型保存与加载
本章代码:
- https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson7/model_save.py
- https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson7/model_load.py
- https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson7/checkpoint_resume.py
- https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson7/save_checkpoint.py
这篇文章主要介绍了序列化与反序列化,以及 PyTorch 中的模型保存于加载的两种方式,模型的断点续训练。
序列化与反序列化
模型在内存中是以对象的逻辑结构保存的,但是在硬盘中是以二进制流的方式保存的。
序列化是指将内存中的数据以二进制序列的方式保存到硬盘中。PyTorch 的模型保存就是序列化。
反序列化是指将硬盘中的二进制序列加载到内存中,得到模型的对象。PyTorch 的模型加载就是反序列化。
PyTorch 中的模型保存与加载
torch.save
torch.save(obj, f, pickle_module, pickle_protocol=2, _use_new_zipfile_serialization=False)
主要参数:
- obj:保存的对象,可以是模型。也可以是 dict。因为一般在保存模型时,不仅要保存模型,还需要保存优化器、此时对应的 epoch 等参数。这时就可以用 dict 包装起来。
- f:输出路径
其中模型保存还有两种方式:
保存整个 Module
这种方法比较耗时,保存的文件大
torch.savev(net, path)
只保存模型的参数
推荐这种方法,运行比较快,保存的文件比较小
state_sict = net.state_dict()
torch.savev(state_sict, path)
下面是保存 LeNet 的例子。在网络初始化中,把权值都设置为 2020,然后保存模型。
import torch
import numpy as np
import torch.nn as nn
from common_tools import set_seed
class LeNet2(nn.Module):
def __init__(self, classes):
super(LeNet2, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 6, 5),
nn.ReLU(),
nn.MaxPool2d(2, 2),
nn.Conv2d(6, 16, 5),
nn.ReLU(),
nn.MaxPool2d(2, 2)
)
self.classifier = nn.Sequential(
nn.Linear(16*5*5, 120),
nn.ReLU(),
nn.Linear(120, 84),
nn.ReLU(),
nn.Linear(84, classes)
)
def forward(self, x):
x = self.features(x)
x = x.view(x.size()[0], -1)
x = self.classifier(x)
return x
def initialize(self):
for p in self.parameters():
p.data.fill_(2020)
net = LeNet2(classes=2019)
# "训练"
print("训练前: ", net.features[0].weight[0, ...])
net.initialize()
print("训练后: ", net.features[0].weight[0, ...])
path_model = "./model.pkl"
path_state_dict = "./model_state_dict.pkl"
# 保存整个模型
torch.save(net, path_model)
# 保存模型参数
net_state_dict = net.state_dict()
torch.save(net_state_dict, path_state_dict)
运行完之后,文件夹中生成了``model.pkl和
model_state_dict.pkl`,分别保存了整个网络和网络的参数
torch.load
torch.load(f, map_location=None, pickle_module, **pickle_load_args)
主要参数:
- f:文件路径
- map_location:指定存在 CPU 或者 GPU。
加载模型也有两种方式
加载整个 Module
如果保存的时候,保存的是整个模型,那么加载时就加载整个模型。这种方法不需要事先创建一个模型对象,也不用知道模型的结构,代码如下:
path_model = "./model.pkl"
net_load = torch.load(path_model)
print(net_load)
输出如下:
LeNet2(
(features): Sequential(
(0): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
(1): ReLU()
(2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(3): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(4): ReLU()
(5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(classifier): Sequential(
(0): Linear(in_features=400, out_features=120, bias=True)
(1): ReLU()
(2): Linear(in_features=120, out_features=84, bias=True)
(3): ReLU()
(4): Linear(in_features=84, out_features=2019, bias=True)
)
)
只加载模型的参数
如果保存的时候,保存的是模型的参数,那么加载时就参数。这种方法需要事先创建一个模型对象,再使用模型的load_state_dict()
方法把参数加载到模型中,代码如下:
path_state_dict = "./model_state_dict.pkl"
state_dict_load = torch.load(path_state_dict)
net_new = LeNet2(classes=2019)
print("加载前: ", net_new.features[0].weight[0, ...])
net_new.load_state_dict(state_dict_load)
print("加载后: ", net_new.features[0].weight[0, ...])
模型的断点续训练
在训练过程中,可能由于某种意外原因如断点等导致训练终止,这时需要重新开始训练。断点续练是在训练过程中每隔一定次数的 epoch 就保存模型的参数和优化器的参数,这样如果意外终止训练了,下次就可以重新加载最新的模型参数和优化器的参数,在这个基础上继续训练。
下面的代码中,每隔 5 个 epoch 就保存一次,保存的是一个 dict,包括模型参数、优化器的参数、epoch。然后在 epoch 大于 5 时,就break
模拟训练意外终止。关键代码如下:
if (epoch+1) % checkpoint_interval == 0:
checkpoint = {"model_state_dict": net.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"epoch": epoch}
path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch)
torch.save(checkpoint, path_checkpoint)
在 epoch 大于 5 时,就break
模拟训练意外终止
if epoch > 5:
print("训练意外中断...")
break
断点续训练的恢复代码如下:
path_checkpoint = "./checkpoint_4_epoch.pkl"
checkpoint = torch.load(path_checkpoint)
net.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']
scheduler.last_epoch = start_epoch
需要注意的是,还要设置scheduler.last_epoch
参数为保存的 epoch。模型训练的起始 epoch 也要修改为保存的 epoch。
参考资料
如果你觉得这篇文章对你有帮助,不妨点个赞,让我有更多动力写出好文章。
[PyTorch 学习笔记] 7.1 模型保存与加载的更多相关文章
- 驱动开发学习笔记. 0.07 Uboot链接地址 加载地址 和 链接脚本地址
驱动开发学习笔记. 0.07 Uboot链接地址 加载地址 和 链接脚本地址 最近重新看了乾龙_Heron的<ARM 上电启动及 Uboot 代码分析>(下简称<代码分析>) ...
- tensorflow 模型保存与加载 和TensorFlow serving + grpc + docker项目部署
TensorFlow 模型保存与加载 TensorFlow中总共有两种保存和加载模型的方法.第一种是利用 tf.train.Saver() 来保存,第二种就是利用 SavedModel 来保存模型,接 ...
- tensorflow实现线性回归、以及模型保存与加载
内容:包含tensorflow变量作用域.tensorboard收集.模型保存与加载.自定义命令行参数 1.知识点 """ 1.训练过程: 1.准备好特征和目标值 2.建 ...
- sklearn模型保存与加载
sklearn模型保存与加载 sklearn模型的保存和加载API 线性回归的模型保存加载案例 保存模型 sklearn模型的保存和加载API from sklearn.externals impor ...
- TensorFlow构建卷积神经网络/模型保存与加载/正则化
TensorFlow 官方文档:https://www.tensorflow.org/api_guides/python/math_ops # Arithmetic Operators import ...
- Tensorflow模型保存与加载
在使用Tensorflow时,我们经常要将以训练好的模型保存到本地或者使用别人已训练好的模型,因此,作此笔记记录下来. TensorFlow通过tf.train.Saver类实现神经网络模型的保存和提 ...
- 转 tensorflow模型保存 与 加载
使用tensorflow过程中,训练结束后我们需要用到模型文件.有时候,我们可能也需要用到别人训练好的模型,并在这个基础上再次训练.这时候我们需要掌握如何操作这些模型数据.看完本文,相信你一定会有收获 ...
- TensorFlow的模型保存与加载
import os os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' import tensorflow as tf #tensorboard --logdir=&qu ...
- Entity Framework学习笔记(五)----Linq查询(2)---贪婪加载
请注明转载地址:http://www.cnblogs.com/arhat 在上一章中,我们使用了Linq对Entity Framework进行了一个查询,但是通过学习我们却发现了懒加载给我来的性能上的 ...
随机推荐
- Spring/Springboot——JavaConfig
1.认识JavaConfig JavaConfig是Spring的一个子项目,在Spring4之后成为一个核心功能 JavaConfig中使用的注解: @Configuration 在类上打上这一标签 ...
- 下载的附件名总乱码?你该去读一下 RFC 文档了!
纸上得来终觉浅,绝知此事要躬行 Web 开发过程中,相信大家都遇到过附件下载的场景,其中,各浏览器下载后的文件名中文乱码问题或许一度让你苦恼不已. 网上搜索一下,大部分都是通过Request Head ...
- day1 linux常用命令(一)
- 【算法•日更•第三十期】区间动态规划:洛谷P4170 [CQOI2007]涂色题解
废话不多说,直接上题: P4170 [CQOI2007]涂色 题目描述 假设你有一条长度为5的木版,初始时没有涂过任何颜色.你希望把它的5个单位长度分别涂上红.绿.蓝.绿.红色,用一个长度为5的字符 ...
- IDEA中列编辑
快捷键 :Alt+Shift+insert,也可以按住Alt+Shift时,点击要编辑部分
- SparkStreaming支持的业务场景
目前而言SparkStreaming 主要支持以下三种业务场景: ➢ 无状态操作:只关注当前的DStream中的实时数据.例如:只对当前DStream中的数据做正确性校验.(下一条与上一条数据无关) ...
- Solon详解(五)- Solon扩展机制之Solon Plugin
Solon 中也有一种非常解耦的扩展机制:Solon Plugin.这种扩展机制和Spring Factories很像,和SPI也很像. 一.Solon 中的扩展机制 在Solon的扩展插件加载机制, ...
- Django中views数据查询使用locals()函数进行优化
优化场景 利用视图函数(views)查询数据之后可以通过上下文context.字典.列表等方式将数据传递给HTML模板,由template引擎接收数据并完成解析.但是通过context传递数据可能就存 ...
- Dockerfile文件万字全面解析
阅读目录 目录 阅读目录 用法 格式 Parser directives escape 环境替换 .dockerignore file FROM RUN CMD LABEL MAINTAINER EX ...
- oracle之回顾二
TCL 1. 事务(Transaction) 事务(Transaction)是一个操作序列.这些操作要么都做, 要么都不做, 是一个不可分割的工作单元, 是数据库环境中的最小工作单元. 1.1 事务 ...