11、torchvision.model 现有网络模型的使用和修改

xiao52x

已于 2022-04-16 20:26:50 修改

阅读量5.7k

点赞数 6

分类专栏：轻松学pytorch 文章标签： python

于 2022-04-16 17:44:57 首次发布

本文链接：https://blog.csdn.net/weixin_44953928/article/details/124216322

版权

轻松学pytorch 专栏收录该内容

15 篇文章 3 订阅

订阅专栏

本文介绍了如何在PyTorch中使用和修改VGG16模型。首先展示了如何加载预训练和未预训练的模型，并打印了模型结构。接着，通过两种方式对模型进行改动，以适应CIFAR10数据集的10个类别，一种是添加新的线性层，另一种是直接修改已有层。最后，讨论了模型的保存和加载方法，包括直接保存整个模型和仅保存权重。

摘要由CSDN通过智能技术生成

一、VGG（CNN经典网络模型）

1、VGG16

torchvision.models.vgg16(pretrained: bool = False, progress: bool = True, **kwargs: Any)

pretrained (bool) – If True, returns a model pre-trained on ImageNet(如果为真，则返回在 ImageNet （是数据集）上预训练的模型)
progress (bool) – If True, displays a progress bar of the download to stderr（如果为 True，则显示下载到 stderr 的进度条

）

# 只是加载网络模型
vgg16_false=torchvision.models.vgg16(pretrained=False) 

#从网络中下载模型（训练好的模型）
vgg16_true=torchvision.models.vgg16(pretrained=True)

在这里插入图片描述

对vgg16模型进行改动

CIFAR10数据集是 10个类别
VGG16输出是1000个类别
VGG 加一层输出10个类别

①、模型的使用

import torchvision
# 直接调用，实例化模型，pretrained代表是否下载预先训练好的参数
vgg16_false = torchvision.models.vgg16(pretrained = False)
vgg16_ture = torchvision.models.vgg16(pretrained = True)
print(vgg16_ture)

输出结果：可以看到VGG16的网络结构

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

②- 模型修改

新增

vgg16_ture.classifier.add_module("add_linear",nn.Linear(1000,10)) # 在vgg16的classfier里加一层
print(vgg16_ture)

只看classfier部分，可以看到一个新增的一层

(classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
    (add_linear): Linear(in_features=1000, out_features=10, bias=True)
  )

修改

print(vgg16_false)
vgg16_false.classifier[6] = nn.Linear(4096,10) # 修改对应层,编号相对应
print(vgg16_false)

#修改前：
(classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )

#修改后：
(classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=10, bias=True)
  )

⑤、模型保存和读取

import torch
import torchvision

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

# 保存方法 1 
torch.save(vgg16,"vgg16_method1.pth") # 保存结构模型和参数、保存路径
# 加载模型 1
model = torch.load("vgg16_method1.pth")

# 保存方式 2 -- 以字典方式只保存参数（官方推荐),
torch.save(vgg16.state_dict(),"vgg_method2.pth") 
# 加载方式 2 -- 要恢复网络模型
model = torch.load("vgg_method2.pth")
vgg16 = torchvision.models.vgg16(pretrained = True)
vgg16.load_state_dict(torch.load("vgg_method2.pth"))

方法一保存：直接加载模型即可

# 保存方法 1 
torch.save(vgg16,"vgg16_method1.pth") # 保存结构模型和参数、保存路径
# 加载模型 1
model = torch.load("vgg16_method1.pth")

方法二保存：字典方式只保存参数：j加载模型还的调用原来的神经网络

# 保存方式 2 -- 以字典方式只保存参数（官方推荐),
torch.save(vgg16.state_dict(),"vgg_method2.pth") 
# 加载方式 2 -- 要恢复网络模型
model = torch.load("vgg_method2.pth")
vgg16 = torchvision.models.vgg16(pretrained = True)
vgg16.load_state_dict(torch.load("vgg_method2.pth"))