在 Pytorch 中正确加载预训练的权重

Question

我想通过加载预训练的视觉变换器模型、修改其最后一层并使用我自己的数据对其进行训练来执行迁移学习。

因此，我正在加载数据集，执行类似于 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

？

我加载预训练权重的方式有问题吗？

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