Правильная загрузка предварительно обученных весов в Pytorch
Я хотел бы выполнить трансферное обучение, загрузив предварительно обученную модель преобразователя зрения, изменив ее последний слой и обучив ее на собственных данных.
Следовательно, я загружаю свой набор данных, выполняю типичное преобразование, аналогичное ImageNet, затем загружаю модель, отключаю градацию на всех ее слоях, удаляю последний слой и добавляю обучаемый, используя количество классов моего набора данных. Мой код мог бы выглядеть следующим образом:
#retrained_vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT # requires torchvision >= 0.13, "DEFAULT" means best available
#pretrained_vit = torchvision.models.vit_b_16(weights=pretrained_vit_weights).to(device)
pretrained_vit = torch.hub.load('facebookresearch/deit:main', 'deit_tiny_patch16_224', pretrained=True).to(device)
for parameter in pretrained_vit.parameters():
parameter.requires_grad = False
pretrained_vit.heads = nn.Linear(in_features=192, out_features=len(class_names)).to(device)
optimizer(torch.optim.Adam(params=pretrained_vit.parameters(), ... )
loss_fn = torch.nn.CrossEntropyLoss()
esults = engine.train(model=pretrained_vit, ..., ... )
Когда я использую torchvision.models.ViT_B_16_Weights.DEFAULT, код работает плавно, и я могу без проблем запустить свой код. Однако, когда я вместо этого использую deit_tiny_patch16_224 и устанавливаю requires_grade = False, я получаю следующую ошибку:
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
Когда для переменной установлено значение True, код работает гладко, но обучение действительно плохое, поскольку у меня было очень мало изображений. Как мне правильно установить параметры deit_tiny_patch16_224 на parameter.requires_grad = False?
Есть ли проблема с тем, как я загружаю предварительно подготовленные веса?
Да, я так думаю. Поскольку я просматриваю все параметры модели и отключаю градиенты, я предполагаю, что это действительно nn.Module.requires_grad_
Ответы 1
Если вы посмотрите описание модели, распечатав его, вы увидите полностью связанный слой классификатора в качестве ключевого имени "head", а не "heads". На моей стороне работает следующий код:
for parameter in pretrained_vit.parameters():
parameter.requires_grad = False
pretrained_vit.head = nn.Linear(in_features=192, out_features=10)
pretrained_vit(torch.rand(1,3,224,224)).mean().backward()
Я рекомендую использовать nn.Module.requires_grad_ вместо того, чтобы самостоятельно устанавливать атрибут для каждого параметра тензора. Имейте в виду, что с вашим текущим кодом вся модель будет заморожена, включая слой классификатора, поэтому вы можете захотеть разморозить этот слой:
pretrained_vit.requires_grad_(False)
pretrained_vit.head = nn.Linear(in_features=192, out_features=10)
pretrained_vit.head.requires_grad_(True)
О, я думаю, у двух моделей были разные названия, и это помогло. Спасибо за помощь :). Что касается вашего второго пункта, какая именно разница, если я установлю все, кроме последнего слоя, через цикл? Разве не то же самое и с вашим предложением?
Вместо того, чтобы назначать атрибут requires_grad всех ваших тензоров False, вы можете просто вызвать requires_grad_ родительского nn.Module, это просто удобнее и понятнее, но ваш способ тоже работает (пока вы не затрагиваете последний слой).
Вы используете nn.Module.requires_grad_?