python - 使用预训练的 ImageNet 模型进行 PyTorch 迁移学习

Question

我想在已在 ImageNet 上训练的模型上使用迁移学习来创建图像分类器。

如何用我自己的自定义分类器替换 torchvision.models ImageNet 分类器的最后一层？

Answer 1

获取预训练的 ImageNet 模型(resnet152 具有最佳准确度):

from torchvision import models
# https://pytorch.org/docs/stable/torchvision/models.html
model = models.resnet152(pretrained=True)

打印出它的结构，以便我们可以与最终状态进行比较:

print(model)

从模型中删除最后一个模块(通常是单个全连接层):

classifier_name, old_classifier = model._modules.popitem()

卡住模型的特征检测器部分的参数，以便它们不会通过反向传播进行调整:

for param in model.parameters():
    param.requires_grad = False

创建一个新的分类器:

classifier_input_size = old_classifier.in_features

classifier = nn.Sequential(OrderedDict([
                           ('fc1', nn.Linear(classifier_input_size, hidden_layer_size)),
                           ('activation', nn.SELU()),
                           ('dropout', nn.Dropout(p=0.5)),
                           ('fc2', nn.Linear(hidden_layer_size, output_layer_size)),
                           ('output', nn.LogSoftmax(dim=1))
                           ]))

我们的分类器的模块名称需要与删除的模块名称相同。将我们的新分类器添加到特征检测器的末尾:

model.add_module(classifier_name, classifier)

最后，打印出新网络的结构:

print(model)