python - 我需要更改哪些参数才能从头开始训练 pytorch 模型？

Question

我按照本教程训练了一个用于实例分割的 pytorch 模型: https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html

我不想在与 COCO 完全无关的完全不同的数据和类别上训练模型。我需要进行哪些更改才能重新训练模型。根据我的阅读，我猜测除了有正确数量的类(class)外，我只需要训练这条线:

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

到

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False)

但我注意到还有另一个参数:pretrained_backbone=True, trainable_backbone_layers=None 是否也应该更改它们？

Answer 1

函数签名是

torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True, trainable_backbone_layers=3, **kwargs)

设置 pretrained=False 将告诉 PyTorch 不要下载在 COCO train2017 上预训练的模型。您想要它，因为您对培训感兴趣。

通常，如果您想在不同的数据集上进行训练，这就足够了。

当您设置 pretrained=False 时，PyTorch 将在 ImageNet 上下载预训练的 ResNet50。默认情况下，它会卡住前两个名为 conv1 和 layer1 的 block 。 Faster R-CNN 论文就是这样做的，它卡住了预训练主干的初始层。

(只需打印模型以检查其结构)。

layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1', 'conv1'][:trainable_layers]

现在，如果您甚至不想卡住前两层，您可以设置 trainable_backbone_layers=5(当您设置 pretrained_backbone=False 时自动完成) ，这将从头开始训练整个 resnet 主干。

检查 PR#2160 .