Pytorch中的保存与加载

path_model = 'model.pkl'
path_state_dict = 'model_state_dict.pkl'
1、保存(序列化)
    # obj:对象
    # f:输出路径
    torch.save(obj,f)
    
2、保存模型参数 # 数据以字典形式存储
    # 获取模型的可学习参数
    state_dict = net.state_dict()
    # 将参数保存
    torch.save(state_dict,f)

3、加载(反序列化)
    # f:文件路径
    @当使用gpu时用map_location
    # map_location:指定存放位置(cpu或gpu)
    net = torch.load(f,map_location)

4、加载模型参数
    # 加载数据
    state_dict_load = torch.load(f,map_location)
    # 将数据加载至模型
    net.load_state_dict(state_dict_load)
    
5、断点续训练/定时保存模型及优化器等训练好的数据
    @网络模型和优化器随学习而不断变化
    # 需要保存的数据
    checkpoint = {
        # 模型数据
        'model_state_dict':net.state_dict(),
        # 优化器数据
        'optimizer_state_dict':optimizer.state_dict(),
        # epoch
        'epoch':epoch
        }
    path = 'checkpoint_{}_epoch.pkl'.format(epoch)
    torch.save(checkpoint,path)
    
    # 加载
    checkpoint = torch.load(path)
    net.load_state_dict(checkpoint[model_state_dict'])
    opimzer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch']

 

  • 0
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值