Pytorch保存训练模型

机器学习与深度学习 专栏收录该内容
3 篇文章 0 订阅

简介

我们在使用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
    点赞
  • 0
    评论
  • 0
    收藏
  • 一键三连
    一键三连
  • 扫一扫,分享海报

相关推荐
©️2020 CSDN 皮肤主题: 深蓝海洋 设计师:CSDN官方博客 返回首页
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值