torch在vgg16预训练模型上添加新的层

gggi520

已于 2024-02-21 23:23:28 修改

阅读量3.3k

点赞数 4

CC 4.0 BY-SA版权

分类专栏： torch 文章标签： python

于 2022-04-04 11:09:25 首次发布

本文链接：https://blog.youkuaiyun.com/gggi520/article/details/123950267

torch 专栏收录该内容

2 篇文章

订阅专栏

本文介绍了如何在PyTorch中加载预训练的VGG16模型，并展示了如何在不同层结构上添加新的层。通过示例代码，展示了如何在VGG、features和classifier后面添加新的线性层。同时，还提供了获取模型中间层输出的方法，以ReLU层为例，展示了如何遍历模型并捕获特定层的输出。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import torch
from torch import nn
import torchvision
vgg16=torchvision.models.vgg16(pretrained=False)
#如果需要预训练模型就让pretrained=True
train_data=torchvision.datasets.CIFAR10('./data',train=True, 
			transform=torchvision.transforms.ToTensor(),download=True)
print(vgg16)

此时输出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)
  )
)

添加新层，首先要注意层结构，这里有3个层结构VGG、features、classifier，需要修改哪个就对应添加哪个

比如在VGG后面加层：

vgg16.add_module('Linear',nn.Linear(1000,10))
#vgg16=VGG

比如在classifier后面加层：

vgg16.classifier.add_module('Linear',nn.Linear(1000,10))

比如在features中间加层：

vgg16.features[28].add_module('Conv2d',nn.Conv2d(512, 512, kernel_size=(3, 3),
							 stride=(1, 1), padding=(1, 1)))

note:PyTorch预训练模型修改、增删特定层
https://blog.youkuaiyun.com/lavinia_chen007/article/details/114833608

#获取中间层输出

import torch
import torch.nn as nn
model = nn.Sequential(
            nn.Conv2d(3, 9, 1, 1, 0, bias=False),
            nn.BatchNorm2d(9),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d((1, 1)),
        )

假如想要获得ReLu的输出

x = torch.rand([2, 3, 224, 224])
for i in range(len(model)):
    x = model[i](x)
    if i == 2:
        ReLu_out = x
print('ReLu_out.shape：\n\t',ReLu_out.shape)
print('x.shape：\n\t',x.shape)