Pytorch：模型的保存与加载

idotc

于 2019-01-21 13:55:37 发布

阅读量514

点赞数 3

分类专栏：深度学习实践文章标签： Pytorch python 深度学习

本文链接：https://blog.youkuaiyun.com/qq_21578849/article/details/86573043

版权

本文介绍了Pytorch中模型保存与加载的两种方法，包括保存整个模型和仅保存模型参数。推荐使用保存参数的方式，因为它更灵活，可用于迁移学习。加载时，torch.load返回OrderedDict，加载过程依赖于key的匹配。当模型结构改变但希望部分加载预训练参数时，可以通过设置strict=False或自定义过滤器来实现。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

模型保存与加载常用有两种方法，第一种是保存整个模型，包括模型的结构和参数；第二种是保存模型的参数。推荐使用第二种，因为模型一旦很大，第一种加载耗时长，其次第二种加载方式更加灵活，可以加载其他模型的预训练参数，从而使用迁移学习的方法减小训练时长。

一、保存/加载整个模型

保存模型：

	torch.save(net, 'model_net1.pkl')

加载模型

	net_parm =  'model_net1.pkl'
    net = torch.load(net_parm)

二、保存/加载模型参数

保存参数：

  	 torch.save({
          'epoch': nums_epoch,
          'state_dict': net.state_dict(),
      }, 'model_net.pkl')

加载参数：

	  cuda_gpu = torch.cuda.is_available()
      if(cuda_gpu):
            net = torch.nn.DataParallel(net, device_ids=gpus).cuda()
      if os.path.isfile(net_parm):
            print(

最低0.47元/天解锁文章