由 torch.save(model.state_dict(), checkpoint_path) 展开的学习

资料仅作为自己复习回顾
参考资料1

torch.save功能
保存模型等相关参数,以及读取保存之后的文件

一、如何保存模型参数、优化器参数等

model = Net() # 网络
optimizer = optim.Adam(model.parameters(), lr=args.lr) # 优化器
# 假设在某个epoch,我们要保存模型参数,优化器参数以及epoch
# 1、先建立一个字典,保存三个参数
state = {‘net’:model.state_dict(), ‘optimizer’:optimizer.state_dict(), ‘epoch’:epoch}
# 2、调用torch.save() 保存模型
torch.save(state, dir)

可以看到,我们的state可以只有模型参数,也可以包含优化器、epoch参数

二、对保存的模型我们要如何读取

当你想恢复某一阶段的训练(或者进行测试)时,那么就可以读取之前保存的网络模型参数等。

checkpoint = torch.load(dir)
model.load_state_dict(checkpoint[‘net’])
optimizer.load_state_dict(checkpoint[‘optimizer’])
start_epoch = checkpoint[‘epoch’] + 1

当前盲区:
model.load_state_dict
optimizer.load_state_dict

  • 2
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值