如何在PyTorch中从ResNet模型中移除最后一个FC层?

45
我正在使用PyTorch中的ResNet152模型。我想从该模型中删除最后一个全连接层。以下是我的代码:

我正在使用PyTorch中的ResNet152模型。我想从该模型中删除最后一个全连接层。以下是我的代码:

from torchvision import datasets, transforms, models
model = models.resnet152(pretrained=True)
print(model)
当我打印模型时,最后几行看起来像这样:<\p>
    (2):  Bottleneck(
      (conv1):  Conv2d(2048,  512,  kernel_size=(1,  1),  stride=(1,  1),  bias=False)
      (bn1):  BatchNorm2d(512,  eps=1e-05,  momentum=0.1,  affine=True,  track_running_stats=True)
      (conv2):  Conv2d(512,  512,  kernel_size=(3,  3),  stride=(1,  1),  padding=(1,  1),  bias=False)
      (bn2):  BatchNorm2d(512,  eps=1e-05,  momentum=0.1,  affine=True,  track_running_stats=True)
      (conv3):  Conv2d(512,  2048,  kernel_size=(1,  1),  stride=(1,  1),  bias=False)
      (bn3):  BatchNorm2d(2048,  eps=1e-05,  momentum=0.1,  affine=True,  track_running_stats=True)
      (relu):  ReLU(inplace)
    )
  )
  (avgpool):  AvgPool2d(kernel_size=7,  stride=1,  padding=0)
  (fc):  Linear(in_features=2048,  out_features=1000,  bias=True)
)

我想从模型中删除最后一个fc层。

我在这里找到了答案(如何将预训练的FC层转换为Pytorch中的CONV层),在那里mexmex似乎提供了我正在寻找的答案:

list(model.modules()) # to inspect the modules of your model
my_model = nn.Sequential(*list(model.modules())[:-1]) # strips off last linear layer

所以我像这样向我的代码添加了这些行:

model = models.resnet152(pretrained=True)
list(model.modules()) # to inspect the modules of your model
my_model = nn.Sequential(*list(model.modules())[:-1]) # strips off last linear layer
print(my_model)

但是这段代码并没有像宣传的那样工作 - 至少对我来说不是。本文的其余部分将详细解释为什么该答案无效,以便此问题不会被关闭为重复。

首先,打印的模型比以前大近5倍。我看到了与以前相同的模型,但后面似乎跟着一个重复的模型,但可能是压平的。

    (2): Bottleneck(
      (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
    )
  )
  (avgpool): AvgPool2d(kernel_size=7, stride=1, padding=0)
  (fc): Linear(in_features=2048, out_features=1000, bias=True)
)
(1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
(2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(3): ReLU(inplace)
(4): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(5): Sequential(
  . . . this goes on for ~1600 more lines . . .
  (415): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (416): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (417): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (418): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (419): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (420): ReLU(inplace)
  (421): AvgPool2d(kernel_size=7, stride=1, padding=0)
)

第二,全连接层仍然存在。在它后面的Conv2D层看起来就像ResNet152的第一层。

第三,如果我尝试调用my_model.forward(),PyTorch会报告一个大小不匹配的错误。它期望大小为[1, 3, 224, 224],但输入是[1, 1000]。因此,看起来整个模型(减去fc层)的副本正在附加到原始模型上。

总之,我在SO上找到的唯一答案实际上并不起作用。


不是很确定,但是基本的删除操作应该可以在这里起作用,del(model['fc'])。你可以试一下吗? - Parthapratim Neog
4个回答

58

对于ResNet模型,你可以使用children属性访问层,因为pytorch中的ResNet模型由nn模块组成。(在pytorch 0.4.1上测试过)

model = models.resnet152(pretrained=True)
newmodel = torch.nn.Sequential(*(list(model.children())[:-1]))
print(newmodel)

更新:尽管并没有一种普遍的答案可以适用于所有pytorch模型的问题,但它应该适用于所有良好结构化的模型。您添加到模型中的已存在的层(例如torch.nn.Lineartorch.nn.Conv2dtorch.nn.BatchNorm2d等)都基于torch.nn.Module类。如果您实现自定义层并将其添加到网络中,则应从pytorch的torch.nn.Module类继承它。正如文档中所写,children属性让您访问类/模型/网络的模块。

def children(self):
        r"""Returns an iterator over immediate children modules.  

更新:需要注意的是,children() 方法返回“即时”的模块。如果您的网络的最后一个模块是一个序列模块,则该方法将返回整个序列模块。


我的回答已更新,附上了文档参考,请提醒我是否有遗漏。 - unlut
2
为了保留层的名称,您可以执行以下操作:newmodel = torch.nn.Sequential(OrderedDict([*(list(model.named_children())[:-1])])) - elbe
使用nn.Sequential时,您实际上是通过一个简单的串联来替换前向函数。那样做会破坏ResNet中的跳跃连接吗? - undefined
@TimKuipers 实际上,如果你打印出模型本身,你会发现它由Bottleneck类的块组成,而不是连续连接的conv relu batchnorms。每个Bottleneck块都是一个残差连接块,你可以在这里查看Bottleneck的前向函数:https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py。最终relu的结果是通过原始输入(x)和中间块输出的求和来计算的。 - undefined

28

您可以通过简单地执行以下步骤来完成:

Model.fc = nn.Sequential()

或者你可以创建身份验证层:

class Identity(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, x):
        return x

使用它替换fc层:

Model.fc = Identity()

11
PyTorch现在有一个内置的identity模块:https://pytorch.org/docs/stable/nn.html#identity - dkv
1
正确的调用方式是torch.nn.Identity()。节省打开链接查找正确调用方法的时间。 - abby37
我爱你!我爱你!这种相同的技术可以用来将内置的PyTorch网络(如ResNet50)从单输出改为多输出! - cdahms

16

如果你不仅想去掉模型的最后一个全连接层,而是想用自己的层替换它,从而利用迁移学习技术,你可以这样做:

import torch.nn as nn
from collections import OrderedDict

n_inputs = model.fc.in_features

# add more layers as required
classifier = nn.Sequential(OrderedDict([
    ('fc1', nn.Linear(n_inputs, 512))
]))

model.fc = classifier

3

来自 PyTorch 教程 “调整 TorchVision 模型”

Here we use Resnet18, as our dataset is small and only has two classes. When we print the model, we see that the last layer is a fully connected layer as shown below:

(fc): Linear(in_features=512, out_features=1000, bias=True)

Thus, we must reinitialize model.fc to be a Linear layer with 512 input features and 2 output features with:

model.fc = nn.Linear(512, num_classes)

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接