基于预训练ImageNet模型的PyTorch迁移学习

获得一个预先训练过的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)

打印(模型)