Pytorch保存训练模型

最新推荐文章于 2024-04-28 19:15:00 发布

牧之原

最新推荐文章于 2024-04-28 19:15:00 发布

阅读量749

点赞数

分类专栏：机器学习与深度学习文章标签：深度学习机器学习 pytorch 神经网络

本文链接：https://blog.csdn.net/Mr_TangGuo/article/details/115469360

版权

机器学习与深度学习专栏收录该内容

4 篇文章 0 订阅

订阅专栏

在使用Pytorch训练模型后，通常需要保存模型以便后续使用。本文介绍了两种方法：一是仅保存模型参数，通过`state_dict()`进行保存和加载；二是保存整个模型，直接使用`torch.save()`和`torch.load()`进行操作。这两种方法能帮助开发者快速地保存和恢复模型状态。

摘要由CSDN通过智能技术生成

简介

我们在使用Pytorch深度学习框架训练自己的模型时，需要保存训练结果以便于下次直接使用训练好的模型对新数据进行预测。本文将介绍Pytorch官方提供的两种模型保存方法

内容

方法1：仅保存模型参数

代码如下：

# 保存模型状态字典，PATH为保存的路径以及文件名
# 例如： PATH='models/myModel.pt'
torch.save(model.state_dict(), PATH)

# 加载模型
# 先实例化一个模型
model = TheModelClass(...)
# 加载状态字典
model.load_state_dict(torch.load(PATH))

方法2：保存整个模型

# 保存
torch.save(model, PATH)
# 加载
model = torch.load(PATH)

牧之原

关注

0
点赞
踩
5

收藏

觉得还不错? 一键收藏
0
评论
Pytorch保存训练模型

简介我们在使用Pytorch深度学习框架训练自己的模型时，需要保存训练结果以便于下次直接使用训练好的模型对新数据进行预测。本文将介绍Pytorch官方提供的两种模型保存方法内容方法1：仅保存模型参数代码如下：# 保存模型状态字典，PATH为保存的路径以及文件名# 例如： PATH='models/myModel.pt'torch.save(model.state_dict(), PATH)# 加载模型# 先实例化一个模型model = TheModelClass(...)# 加载状态
复制链接

扫一扫

专栏目录