pytorch模型的加载与保存

最新推荐文章于 2024-01-11 21:59:31 发布

wsp+

最新推荐文章于 2024-01-11 21:59:31 发布

阅读量197

点赞数

文章标签： pytorch 深度学习 python

本文链接：https://blog.csdn.net/qq_42193790/article/details/128284628

版权

保存

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'))