Pytorch深度学习-网络模型的保存与读取（小土堆）

最新推荐文章于 2024-09-28 21:53:15 发布

Mr chenxizhi

最新推荐文章于 2024-09-28 21:53:15 发布

阅读量359

点赞数 3

文章标签：深度学习 pytorch 人工智能

本文链接：https://blog.csdn.net/weixin_53141706/article/details/140552649

版权

代码示例

import torchvision  
import torch  
  
  
vgg16 = torchvision.models.vgg16(pretrained=False)  
#保存方式1:不仅保存结构，还保存参数  
torch.save(vgg16,"vgg16_model1.pth")  
  
#保存方式1->加载模型  
torch.load("vgg16_medel1.pth")  
  
#保存方式2:不保存结构，只保存参数（官方推荐）  
torch.save(vgg16.state_dict(),"vgg16_model2.pth")  
  
#保存方式2->加载模型  
#因为只保存参数，所以使用时需要重新加载模型  
vgg16 = torchvision.models.vgg16(pretrained=False)  
#加载保存的参数  
model = torch.load("vgg16_model2.pth")  
#将保存的参数读取到模型之中  
vgg16.state_dict(model)