PyTorch transfer learning with pre-trained ImageNet model

问题

I want to create an image classifier using transfer learning on a model already trained on ImageNet.

How do I replace the final layer of a torchvision.models ImageNet classifier with my own custom classifier?

回答1:

Get a pre-trained ImageNet model (resnet152 has the best accuracy):

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

Print out its structure so we can compare to the final state:

print(model)

Remove the last module (generally a single fully connected layer) from model:

classifier_name, old_classifier = model._modules.popitem()

Freeze the parameters of the feature detector part of the model so that they are not adjusted by back-propagation:

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

Create a new classifier:

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))
                           ]))

The module name for our classifier needs to be the same as the one which was removed. Add our new classifier to the end of the feature detector:

model.add_module(classifier_name, classifier)

Finally, print out the structure of the new network:

print(model)

来源：https://stackoverflow.com/questions/53884692/pytorch-transfer-learning-with-pre-trained-imagenet-model