pytorch模型导入、修改、保存、读取

最新推荐文章于 2024-07-25 14:35:42 发布

米饭�

最新推荐文章于 2024-07-25 14:35:42 发布

阅读量658

点赞数

分类专栏： pytorch 文章标签： pytorch python 深度学习

本文链接：https://blog.csdn.net/weixin_43811423/article/details/128258396

版权

pytorch 专栏收录该内容

8 篇文章 0 订阅

订阅专栏

pytorch对模型的操作

导入

这里以vgg11的classifier层为例
导入模型：

vgg11 = torchvision.models.vgg11(pretrained=False)
print(vgg11)

其中classifier的输出如下
在这里插入图片描述

修改

主要的修改方式：

# 在某层中添加层
vgg11.classifier.add_module('new_linear', nn.Linear(1000, 10))
# 修改某层
vgg11.classifier[6] = nn.Linear(4096, 10)
print(vgg11)

修改后：
在这里插入图片描述

保存与读取

# 模型的保存与读取
# 方式1，保存模型和参数
torch.save(vgg11, 'vgg11_method1.pth')
# 在读取时需要保证原模型已经引入
model = torch.load('vgg11_method1.pth')

# 方式2，只保存模型参数，一个字典形式(官方推荐)
torch.save(vgg11.state_dict(), 'vgg11_method2.pth')

vgg11_new = torchvision.models.vgg11(pretrained=False)
vgg11_new.load_state_dict(torch.load('vgg11_method2.pth'))