pytorch模型的加载与保存

保存

import torch
import torchvision

vgg16=torchvision.models.vgg16(pretrained=False)

#模型的保存与加载
#保存方式一
#不光保存模型结构 也保存参数
torch.save(vgg16,'vgg16_model.pth')
#保存方式二
#保存模型参数(官方推荐) 字典形式
torch.save(vgg16.state_dict(),'vgg16_model2.pth')

加载

import torch
import torchvision

#加载模型

#方式一---->保存方式一
model=torch.load('vgg16_model.pth')


#方式二加载模型
vgg16=torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load('vgg16_model2.pth'))