如何改变不同的火炬模型最后一层进行统一的微调?

Question

我有一个torch model，它接收一个预训练模型，并改变预训练模型的最后一个module进行微调。但现实情况是，这些预训练模型的最后一个模块有不同的名称，例如，torchvision.models.resnet*()是fc,timm.models.vit*()是head，其他一些可能包含几个输出层组成最后一个模块。

我想要的是我可以用统一的方式改变这些不同的层。如

last_model = get_last_module(pretrained_model)
last_model = nn.Sequential(...)

。但是这样,pretrained_model不改变了。

例如，我做了一件愚蠢的事:

def __init__(self, base_encoder, pred_dim=512):
"""
dim: feature dimension (default: 2048)
pred_dim: hidden dimension of the predictor (default: 512)
"""
super().__init__()
# create the encoder
# num_classes is the output fc dimension, zero-initialize last BNs
self.encoder = base_encoder
if isinstance(self.encoder, torchvision.models.ResNet):
dim = self.encoder.fc.out_features
# build a 3-layer projector
prev_dim = self.encoder.fc.in_features
self.encoder.fc = nn.Sequential(
nn.Linear(prev_dim, prev_dim, bias=False),
nn.BatchNorm1d(prev_dim),
nn.ReLU(inplace=True),  # first layer
nn.Linear(prev_dim, prev_dim, bias=False),
nn.BatchNorm1d(prev_dim),
nn.ReLU(inplace=True),  # second layer
# nn.Linear(prev_dim, dim),
self.encoder.fc,
nn.BatchNorm1d(dim, affine=False),
)  # output layer
self.encoder.fc[
6
].bias.requires_grad = False  # hack: not use bias as it is followed by BN
elif isinstance(self.encoder, timm.models.VisionTransformer):
dim = self.encoder.head.out_features
# build a 3-layer projector
prev_dim = self.encoder.head.in_features
self.encoder.head = nn.Sequential(
nn.Linear(prev_dim, prev_dim, bias=False),
nn.BatchNorm1d(prev_dim),
nn.ReLU(inplace=True),  # first layer
nn.Linear(prev_dim, prev_dim, bias=False),
nn.BatchNorm1d(prev_dim),
nn.ReLU(inplace=True),  # second layer
# nn.Linear(prev_dim, dim),
self.encoder.head,
nn.BatchNorm1d(dim, affine=False),
)  # output layer
self.encoder.head[
6
].bias.requires_grad = False  # hack: not use bias as it is followed by BN
# build a 2-layer predictor
self.predictor = nn.Sequential(
nn.Linear(dim, pred_dim, bias=False),
nn.BatchNorm1d(pred_dim),
nn.ReLU(inplace=True),  # hidden layer
nn.Linear(pred_dim, dim),
)

Answer 1

使用儿童如何?

pretrained_model_except_last_layer = list(pretrained_model.children())[:-1]
new_model = nn.Sequential(*pretrained_model_except_last_layer, your_new_classfier)