pytorch保存模型、加载模型的两种方式

第一种 保存整个模型文件 

  

   

    

     

    
    

     

      # 保存和加载整个模型
     
    
   

    

     

    
    

     

      torch.save(model_object, 'model.pth')
     
    
   

    

     

    
    

     

      model = torch.load('model.pth')
     
    
  
 

第二种 仅保存模型的权重 

  

   

    

     

    
    

     

      # 仅保存和加载模型参数(推荐使用)
     
    
   

    

     

    
    

     

      torch.save(model_object.state_dict(), 'weights.pth')
     
    
   

    

     

    
    

     

      model.load_state_dict(torch.load('weights.pth'))
     
    
  
 

文章来源: wanghao.blog.csdn.net，作者：AI浩，版权归原作者所有，如需转载，请联系作者。

原文链接：wanghao.blog.csdn.net/article/details/110433626