0

我正在尝试将迁移学习用于图像分割任务,我的计划是使用预训练模型(例如 VGG16)的前几层作为编码器,然后添加我自己的解码器。

所以,我可以加载模型并通过打印它来查看结构:

model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True)
print(model)

我得到这样的:

ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (1): BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  .....
  .....
  .....

例如,我还可以访问特定的图层model.layer3。现在,我正在为某些事情而苦苦挣扎。

  1. 如何切割模型并将每个模块从任何层的开始到结束(例如model.layer3)?
  2. 如何freeze只保留这个被剥离的部分,并保持新添加的模块可用于培训?
4

3 回答 3

1

对于 1):初始化您的 ResNetLightningModule并将其切片直到您需要的部分。然后在此之后添加您自己的头部,并forward按照您需要的顺序进行定义。请参阅此示例,基于迁移学习文档

import torchvision.models as models

class ImagenetTransferLearning(LightningModule):
    def __init__(self):
        super().__init__()

        # init a pretrained resnet
        backbone_tmp = models.resnet50(pretrained=True)
        num_filters = backbone_tmp.fc.in_features
        layers = list(backbone_tmp.children())[:-1]
        self.backbone = nn.Sequential(*layers)

        # use the pretrained model to classify cifar-10 (10 image classes)
        num_target_classes = 10
        self.classifier = nn.Linear(num_filters, num_target_classes)

对于 2):将 BackboneFinetuning 回调传递给您的trainer. 这要求您LightningModule有一个self.backbone包含要冻结的模块的属性,如上面的代码片段所示。如果您需要不同的冻结-解冻行为,您也可以使用BaseFinetuning 回调。

from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import BackboneFinetuning
multiplicative = lambda epoch: 1.5
backbone_finetuning = BackboneFinetuning(200, multiplicative)
trainer = Trainer(callbacks=[backbone_finetuning])
于 2021-05-27T14:01:09.290 回答
1

以下情况适用于 的任何子模块model,但我将在这里回答您的问题model.layer3

  1. model.layer3将为您提供nn.Module与模型的第 3 层相关联的信息。您可以像使用一样直接调用它model

    >>> z = model.layer3(torch.rand(16, 128, 10, 10))
    >>> z.shape
    torch.Size([16, 256, 5, 5])
    
  2. 冻结模型:

    • 您可以将该层置于eval模式,该模式禁用 dropouts 并使 BN 层在训练期间使用统计学习。这是用model.layer3.eval()

    • requires_grad您必须通过切换标志来禁用该层的训练:model.layer3.requires_grad_(False)这将影响所有子参数。

于 2021-02-01T22:01:12.493 回答
0

您可以使用以下方法冻结图层:

pretrained_model.freeze()

https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.core.lightning.html?highlight=Freeze#pytorch_lightning.core.lightning.LightningModule.freeze

于 2021-02-22T03:07:52.550 回答