预训练ViT模型参数冻结：仅解冻新增Adapter层的实现方法

我基于PyTorch Vision的预训练ViT模型作为编码器，给其EncoderBlock结构新增了Adapter模块并完成集成。现在需要冻结ViT的全部原有参数，仅让所有新增的Adapter模块参数可训练，请问该如何操作？

（原EncoderBlock、Adapter以及集成后的代码如下）

原EncoderBlock结构 #

class EncoderBlock(nn.Module):
"""Transformer encoder block."""

def __init__(
    self,
    num_heads: int,
    hidden_dim: int,
    mlp_dim: int,
    dropout: float,
    attention_dropout: float,
    norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
):
    super().__init__()
    self.num_heads = num_heads

    # Attention block
    self.ln_1 = norm_layer(hidden_dim)
    self.self_attention = nn.MultiheadAttention(hidden_dim, num_heads, dropout=attention_dropout, batch_first=True)
    self.dropout = nn.Dropout(dropout)

    # MLP block
    self.ln_2 = norm_layer(hidden_dim)
    self.mlp = MLPBlock(hidden_dim, mlp_dim, dropout)

def forward(self, input: torch.Tensor):
    torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
    x = self.ln_1(input)
    x, _ = self.self_attention(x, x, x, need_weights=False)
    x = self.dropout(x)
    x = x + input

    y = self.ln_2(x)
    y = self.mlp(y)
    return x + y

Adapter模块定义 #

class Adapter(nn.Module):
    def __init__(self, D_features, mlp_ratio=0.25, act_layer=nn.GELU):
        super().__init__()
        D_hidden_features = int(D_features * mlp_ratio)
        self.act = act_layer()
        self.D_fc1 = nn.Linear(D_features, D_hidden_features)
        self.D_fc2 = nn.Linear(D_hidden_features, D_features)
        
    def forward(self, x):
        # x is (BT, HW+1, D)
        xs = self.D_fc1(x)
        xs = self.act(xs)
        xs = self.D_fc2(xs)
        
        x = xs
        return x

集成后的EncoderBlock #

def __init__(
    self,
    num_heads: int,
    hidden_dim: int,
    mlp_dim: int,
    dropout: float,
    attention_dropout: float,
    norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
):
    super().__init__()
    self.num_heads = num_heads

    # Attention block
    self.ln_1 = norm_layer(hidden_dim)
    self.self_attention = nn.MultiheadAttention(hidden_dim, num_heads, dropout=attention_dropout, batch_first=True)
    self.dropout = nn.Dropout(dropout)

    # MLP block
    self.ln_2 = norm_layer(hidden_dim)
    self.mlp = MLPBlock(hidden_dim, mlp_dim, dropout)
    
    self.adapter = Adapter(D_features = hidden_dim)

def forward(self, input: torch.Tensor):
    torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
    x = self.ln_1(input)
    x, _ = self.self_attention(query=x, key=x, value=x, need_weights=False)
    
    x = self.adapter(x)      # 新增的Adapter调用
 
    x = self.dropout(x)

    x = x + input

    y = self.ln_2(x)
    y = self.mlp(y)
    return x + y

可以通过以下步骤实现仅解冻所有Adapter模块的参数：

1. 先冻结整个ViT模型的所有参数
   先把模型的所有参数都设置为不可训练，确保原有预训练参数不会被更新：

   # 假设你的预训练ViT模型实例为vit_model
   for param in vit_model.parameters():
       param.requires_grad = False
   

2. 遍历所有EncoderBlock，解冻其中的Adapter参数
   PyTorch Vision的ViT模型中，编码器层通常存储在vit_model.encoder.layers中，每个元素都是你修改后的带Adapter的EncoderBlock。遍历每个层并开启其adapter模块的参数可训练属性：

   for layer in vit_model.encoder.layers:
       # 遍历Adapter模块下的所有参数
       for param in layer.adapter.parameters():
           param.requires_grad = True
   

3. （可选）验证可训练参数
   可以打印出所有可训练的参数名称，确认只有Adapter的参数被解冻：

   trainable_params = [name for name, param in vit_model.named_parameters() if param.requires_grad]
   print("可训练参数列表：")
   for name in trainable_params:
       print(name)
   

   输出应该类似：

   encoder.layers.0.adapter.D_fc1.weight
   encoder.layers.0.adapter.D_fc1.bias
   encoder.layers.0.adapter.D_fc2.weight
   encoder.layers.0.adapter.D_fc2.bias
   encoder.layers.1.adapter.D_fc1.weight
   ...
   

4. 优化器配置
   在定义优化器时，只传入可训练的参数，提升训练效率：

   optimizer = torch.optim.AdamW(
       [param for param in vit_model.parameters() if param.requires_grad],
       lr=1e-4
   )
   

内容的提问来源于stack exchange，提问作者dtr43