Передача обучения PyTorch с предварительно обученной моделью ImageNet

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)