如何修改冻结的预训练模型的网络结构

在冻结预训练模型的部分参数或层之后，仍可以对模型的结构进行修改，以便更好地适应下游任务。以下是几种常见的修改方法：

1. 添加新的层

• 直接在模型上添加新的层，常用于增加下游任务的分类头或回归头。例如，添加一个新的全连接层用于分类。
• 新增的层会随机初始化，通常不会冻结它们，以便在训练时适应任务需求。

import torch.nn as nn

class ModifiedModel(nn.Module):
    def __init__(self, pretrained_model):
        super(ModifiedModel, self).__init__()
        self.base_model = pretrained_model
        self.new_head = nn.Sequential(
            nn.Linear(pretrained_model.hidden_size, 512),
            nn.ReLU(),
            nn.Linear(512, num_classes)  # num_classes为你的下游任务分类类别数
        )

    def forward(self, x):
        features = self.base_model(x)
        output = self.new_head(features)
        return output

2. 替换特定层

• 有时我们希望替换冻结模型中的某些层，例如替换预训练模型中的分类层以适应新的类别数。
• 可以直接替换 pretrained_model.layer 中的特定层。

# 替换某个层
model.base_model.layer_name = NewLayerClass()  # 替换原有的层

3. 增加跳接连接（Skip Connections）

• 如果希望在不影响冻结层的前提下添加更多上下游层的交互，可以引入跳接连接。例如，将模型的一部分输出送入其他新层，并在后续层进行融合。

class SkipConnectedModel(nn.Module):
    def __init__(self, pretrained_model):
        super(SkipConnectedModel, self).__init__()
        self.base_model = pretrained_model
        self.additional_layer = nn.Linear(pretrained_model.hidden_size, 512)
        self.final_layer = nn.Linear(pretrained_model.hidden_size + 512, num_classes)

    def forward(self, x):
        base_features = self.base_model(x)
        additional_features = self.additional_layer(base_features)
        combined_features = torch.cat([base_features, additional_features], dim=1)
        output = self.final_layer(combined_features)
        return output

4. 解冻部分层

• 虽然整体模型被冻结，但可以选择性地解冻部分层或块进行训练。
• 这通常应用于微调的最后几层，使模型能更好地适应新的数据分布。

# 解冻最后几层
for name, param in model.named_parameters():
    if "layer_name" in name:  # 通过名称选择性解冻
        param.requires_grad = True

5. 修改模型的输入或输出结构

• 如果输入结构不匹配，可以通过添加输入处理层对输入进行预处理。
• 对输出进行修改也类似，可以在最后添加一个处理层来对特征进行适配。

实例代码

假设我们使用的是一个预训练的视觉Transformer模型，可以通过以下代码添加新的任务头，并选择性地解冻模型的后几层：

import torch
import torch.nn as nn

class ModifiedViLT(nn.Module):
    def __init__(self, pretrained_model, num_classes):
        super(ModifiedViLT, self).__init__()
        self.base_model = pretrained_model
        self.new_head = nn.Linear(pretrained_model.hidden_size, num_classes)
        
        # 冻结所有参数
        for param in self.base_model.parameters():
            param.requires_grad = False
            
        # 选择性解冻最后一个Transformer层
        for param in self.base_model.transformer.layers[-1].parameters():
            param.requires_grad = True

    def forward(self, x):
        base_features = self.base_model(x)
        output = self.new_head(base_features)
        return output

以上方法可以帮助在不影响冻结层的前提下，修改预训练模型的结构以更好地适应下游任务。