深度学习——09模型的保存:torch.save()、加载:torch.load()

两种方式

保存模型主要分为两类：
1、保存整个模型
2、保存模型参数

1、第一种

结构模型+模型参数

保存整个网络模型，加载整个网络模型（可能比较耗时）

# 保存方式1
torch.save(vgg16, "vgg16_model1.pth")
# 对应保存的方式1
model = torch.load("vgg16_model1.pth")
print(model)

2、第二种

只保存加载模型参数（推荐）

保存模型的权重参数（速度快，占内存少）

# 保存方式2
torch.save(vgg16.state_dict(),"vgg16_model2.pth")
# 对应保存的方式2
vgg16 = torchvision.models.vgg16(pretrained=False)
model2 = torch.load("vgg16_model2.pth")
vgg16.load_state_dict(model2)

假设网络为：
model = Net(), optimizer = optim.Adam(model.parameters(), lr=args.lr),
假设在某个epoch，要保存模型参数，优化器参数以及epoch
先建立一个字典，保存三个参数：

state = { ‘net’: model.state_dict(), ‘optimizer’: optimizer.state_dict(), ‘epoch’: epoch}

torch.save(state, "./project/mymodel.pth")

当想恢复某一阶段的训练时，那就可以读取之前保存的网络模型参数等。

mymodel= torch.load("./project/mymodel.pth")
model.load_state_dict(mymodel['net'])  #  加载之前的网络模型参数
optimizer.load_state_dict(mymodel['optimizer'])  # 加载之前的优化器的参数
start_epoch = mymodel['epoch'] + 1  #  加载新的训练回合数

深度学习——09模型的保存:torch.save()、加载:torch.load()相关推荐

Tensorflow【实战Google深度学习框架】TensorFlow模型的保存与恢复加载
我们使用TensorFlow进行模型的训练,训练好的模型需要保存,预测阶段我们需要将模型进行加载还原使用,这就涉及TensorFlow模型的保存与恢复加载. 总结一下Tensorflow常用的模型保存 ...
深度学习修炼（二）——数据集的加载
文章目录致谢 2 数据集的加载 2.1 框架数据集的加载 2.2 自定义数据集 2.3 准备数据以进行数据加载器训练致谢 Pytorch自带数据集介绍_godblesstao的博客-CSDN博客_ ...
Pytorch：模型的保存与加载 torch.save()、torch.load()、torch.nn.Module.load_state_dict()
Pytorch 保存和加载模型后缀:.pt 和.pth 1 torch.save() [source] 保存一个序列化(serialized)的目标到磁盘.函数使用了Python的pickle程序用于 ...
Tensorflow深度学习实战之(五)--保存与恢复模型
文章目录一.保存模型二.恢复模型三.使用模型预测一.保存模型在训练完Tensorflow模型为了方便对新的数据进行预测需要保存该模型,Tensorflow提供 tf.train.Saver( ...
深度学习实战——模型推理优化（模型压缩与加速）
忆如完整项目/代码详见github:https://github.com/yiru1225(转载标明出处勿白嫖 star for projects thanks) 目录系列文章目录一.实验思路综 ...
【深度学习中模型评价指标汇总（混淆矩阵、recall、precision、F1、AUC面积、ROC曲线、ErrorRate）】
深度学习中模型好坏的所有评价指标汇总(混淆矩阵.recall.precision.F1score.AUC面积.ROC曲线.ErrorRate) 导航 0.混淆矩阵 1.AUC面积 2.ROC曲线 3. ...
浅谈深度学习:如何计算模型以及中间变量的显存占用大小
原文链接:https://oldpan.me/archives/how-to-calculate-gpu-memory 前言亲,显存炸了,你的显卡快冒烟了! torch.FatalError: cu ...
NVIDIA GPUs上深度学习推荐模型的优化
NVIDIA GPUs上深度学习推荐模型的优化 Optimizing the Deep Learning Recommendation Model on NVIDIA GPUs 推荐系统帮助人在成倍增 ...
【深度学习】模型训练过程可视化思路（可视化工具TensorBoard）
[深度学习]模型训练过程可视化思路(可视化工具TensorBoard) 文章目录 1 TensorBoard的工作原理 2 TensorFlow中生成log文件 3 启动TensorBoard,读取l ...

深度学习——09模型的保存:torch.save()、加载:torch.load()

两种方式

结构模型+模型参数

只保存加载模型参数（推荐）

深度学习——09模型的保存:torch.save()、加载:torch.load()相关推荐

最新文章

热门文章