Zhiyi12 t adapter

src/viv1t/model/core/vivit.py

@@ -16,6 +16,41 @@
 
 FF_ACTIVATIONS = Literal["gelu", "swiglu"]
 
+class SAdapter2(nn.Module):
+    def __init__(self, D_features, mlp_ratio=0.75, act_layer=nn.GELU, skip_connect=True):
+        super().__init__()
+        D_hidden_features = int(D_features * mlp_ratio)
+        self.D_fc1 = nn.Linear(D_features, D_hidden_features)
+        self.D_fc2 = nn.Linear(D_hidden_features, D_features)
+        self.act = act_layer()
+        self.skip_connect = skip_connect  # 添加 skip_connect 参数
+
+    def forward(self, x):
+        xs = self.D_fc1(x)
+        xs = self.act(xs)
+        xs = self.D_fc2(xs)
+        if self.skip_connect:
+            x = x + xs
+        else:
+            x = xs
+        return x
+
+class TAdapter(nn.Module):
+    def __init__(self, D_features, mlp_ratio=0.75, act_layer=nn.GELU):
+        super().__init__()
+        D_hidden_features = int(D_features * mlp_ratio)
+        self.D_fc1 = nn.Linear(D_features, D_hidden_features)
+        self.D_fc2 = nn.Linear(D_hidden_features, D_features)
+        self.act = act_layer()
+
+
+    def forward(self, x):
+        xs = self.D_fc1(x)
+        xs = self.act(xs)
+        xs = self.D_fc2(xs)
+        x = xs
+        # x = rearrange(x, 'B N T D -> (B T) N D')
+        return x
 
 def find_shape(num_patches: int):
     dim1 = math.ceil(math.sqrt(num_patches))
@@ -24,7 +59,6 @@ def find_shape(num_patches: int):
     dim2 = num_patches // dim1
     return dim1, dim2
 
-
 class SinusoidalPositionalEncoding(nn.Module):
     def __init__(
         self,
@@ -67,7 +101,6 @@ def forward(self, inputs: torch.Tensor):
                 outputs += self.pos_encoding[:, :, : inputs.size(2), :]
         return self.dropout(outputs)
 
-
 class Unfold3d(nn.Module):
     def __init__(
         self,
@@ -94,7 +127,6 @@ def forward(self, inputs: torch.Tensor):
         )
         return outputs
 
-
 class UnfoldConv3d(nn.Module):
     def __init__(
         self,
@@ -273,6 +305,7 @@ def unfold_size(dim: int, patch_size: int, padding: int = 0, stride: int = 1):
         return math.floor(((dim + 2 * padding - patch_size) / stride) + 1)
 
     def forward(self, inputs: torch.Tensor):
+        # inputs = inputs.permute(0, 1, 3, 2) # edit input
         outputs = self.tokenizer(inputs)
         _, t, p, _ = outputs.shape
         match self.pos_encoding:
@@ -292,7 +325,6 @@ def forward(self, inputs: torch.Tensor):
                 outputs = self.temporal_pos_encoding(outputs)
         return outputs
 
-
 class PositionalEncodingGenerator(nn.Module):
     """Position Encoding Generator from https://arxiv.org/abs/2102.10882"""
 
@@ -340,12 +372,13 @@ def forward(self, inputs: torch.Tensor):
                 outputs = outputs + self.pos_embedding(outputs)
                 outputs = rearrange(outputs, "b c h w -> b (h w) c")
             case "temporal":
-                outputs = rearrange(outputs, "b t c -> b c t ")
+                # 假设新的输入格式为 (B, N, T, C)
+                b, n, t, d = inputs.shape
+                outputs = rearrange(outputs, "b n t d -> b d (n t)")
                 outputs = outputs + self.pos_embedding(outputs)
-                outputs = rearrange(outputs, "b c t -> b t c")
+                outputs = rearrange(outputs, "b d (n t) -> b n t d", n=n)
         return outputs
 
-
 class Attention(nn.Module):
     def __init__(
         self,
@@ -476,8 +509,7 @@ def forward(
                 outputs = torch.cat((behaviors, pupil_center), dim=-1)
         outputs = self.models[mouse_id](outputs)
         return outputs
-
-
+
 class TransformerBlock(nn.Module):
     def __init__(
         self,
@@ -497,7 +529,8 @@ def __init__(
         flash_attention: bool = True,
         normalize_qk: bool = False,
         grad_checkpointing: bool = False,
-        ff_activation: FF_ACTIVATIONS = "gelu",
+        ff_activation: str = "gelu",
+        mlp_ratio: float = 4.0  # 添加 mlp_ratio 参数，设置默认值为 4.0
     ):
         super(TransformerBlock, self).__init__()
         assert dimension in ("spatial", "temporal")
@@ -512,6 +545,12 @@ def __init__(
             normalize_qk=normalize_qk,
             grad_checkpointing=grad_checkpointing,
         )
+        if dimension == "spatial":
+            self.adapter = SAdapter2(D_features=emb_dim, mlp_ratio=0.5)
+        else:
+            self.adapter = TAdapter(D_features=emb_dim, mlp_ratio=0.5)
+        # self.sadapter2 = SAdapter2(D_features=emb_dim, mlp_ratio=0.25)
+        # self.tadapter2 = TAdapter(D_features=emb_dim, mlp_ratio=0.45)
         match ff_activation:
             case "gelu":
                 mlp_out = mlp_dim
@@ -528,7 +567,7 @@ def __init__(
             nn.Linear(in_features=emb_dim, out_features=mlp_out),
             activation,
             nn.Dropout(p=ff_dropout),
-            nn.Linear(in_features=mlp_dim, out_features=emb_dim),
+            nn.Linear(in_features=mlp_out, out_features=emb_dim),  # 修正这里的 in_features
             nn.Dropout(p=ff_dropout),
         )
         self.pos_encoding = None
@@ -575,15 +614,37 @@ def forward(
                 case "spatial":
                     states = repeat(states, "b d -> b 1 d")
                 case "temporal":
-                    states = repeat(states, "b t d -> b p t d", p=self.spatial_dim)
-                    states = rearrange(states, "b p t d -> (b p) t d")
+                    states = repeat(states, "b d -> b t d", t=inputs.size(1))  # 修正这里的维度
             outputs = outputs + states
-        outputs = self.drop_path1(self.attention(outputs)) + outputs
+        # LayerNorm --> S-MSA ---> Adapter ---> + ---> LayerNorm ---> MLP
+        #    |                                  |
+        #    |--------------------------------->|
+        # outputs = self.drop_path1(self.attention(outputs)) + outputs
+        # outputs = self.drop_path1(adapter_result) + outputs
+        # outputs = self.drop_path2(self.ff(outputs)) + outputs
+        attn_output = self.attention(outputs)
+        adapter_output = self.adapter(attn_output)
+        # sadapter_output = self.sadapter2(attn_output) # Applying SAdapter2 right after attention
+        # tadapter_output = self.tadapter2(attn_output)
+        # adapter_output = sadapter_output + tadapter_output
+        # adapter_output = torch.cat((sadapter_output, tadapter_output), dim=0)
+        outputs = self.drop_path1(adapter_output) + outputs
         outputs = self.drop_path2(self.ff(outputs)) + outputs
+
+        # outputs = self.drop_path1(self.attention(outputs)) + outputs
+        # outputs = self.sadapter2(outputs)  # Applying SAdapter2 after attention
+        # outputs = self.drop_path2(self.ff(outputs)) + outputs
         if self.pos_encoding is not None:
             outputs = self.pos_encoding(outputs)
         return outputs
 
+    def freeze_except_adapter(self):
+        for name, param in self.named_parameters():
+            if 'adapter' in name:
+                param.requires_grad = True
+            else:
+                print(f"{name} is frozen\n");
+                param.requires_grad = False
 
 class ParallelTransformerBlock(nn.Module):
     def __init__(
@@ -658,6 +719,7 @@ def __init__(
             self.spatial_dim = spatial_dim
 
         self.apply(self.init_weight)
+        self.sadapter2 = SAdapter2(D_features=emb_dim, mlp_ratio=0.75)
 
     @staticmethod
     def init_weight(m: nn.Module):
@@ -697,7 +759,8 @@ def parallel_attention(self, inputs: torch.Tensor):
             + self.drop_path1(self.attn_out(attn))
             + self.drop_path2(self.ff_out(ff))
         )
-        return outputs
+        return self.sadapter2(outputs)
+        # return outputs
 
     def forward(
         self,
@@ -726,7 +789,14 @@ def forward(
         else:
             outputs = self.parallel_attention(outputs)
         return outputs
-
+
+    def freeze_except_adapter(self):
+        for name, param in self.named_parameters():
+            if 'sadapter2' in name:
+                param.requires_grad = True
+            else:
+                print(f"{name} is frozen\n");
+                param.requires_grad = False
 
 class Transformer(nn.Module):
     def __init__(
@@ -778,6 +848,9 @@ def __init__(
                 for i in range(depth)
             ]
         )
+    def freeze_blocks(self):
+        for blk in self.blocks:
+            blk.freeze_except_adapter()
 
     def forward(
         self,
@@ -796,7 +869,6 @@ def forward(
             )
         return outputs
 
-
 class ViViT(nn.Module):
     def __init__(
         self,
@@ -894,7 +966,8 @@ def __init__(
             grad_checkpointing=args.grad_checkpointing,
             ff_activation=args.core_ff_activation,
         )
-
+        self.spatial_transformer.freeze_blocks()
+        self.temporal_transformer.freeze_blocks()
         # calculate latent height and width based on num_patches
         new_h, new_w = find_shape(self.tokenizer.output_shape[1])
         self.rearrange = Rearrange("b t (h w) c -> b c t h w", h=new_h, w=new_w)
@@ -909,6 +982,13 @@ def temporal_unfold(self, inputs: torch.Tensor):
         outputs = rearrange(outputs, "b d t pt -> b t (pt d)")
         return outputs
 
+    def freeze_non_transformer_layers(self):
+        for name, param in self.named_parameters():
+            if 'transformer' in name:
+                param.requires_grad = True
+            else:
+                print(f"{name} is frozen in Vivit")
+                param.requires_grad = False
     def forward(
         self,
         inputs: torch.Tensor,
@@ -974,8 +1054,7 @@ def forward(
 
         return outputs
 
-
-@register("vivit")
+@register("vit")
 class ViViTCore(Core):
     def __init__(self, args: Any, input_shape: Tuple[int, int, int, int]):
         """
@@ -1004,6 +1083,7 @@ def __init__(self, args: Any, input_shape: Tuple[int, int, int, int]):
             case _:
                 raise NotImplementedError(f"--behavior_mode {self.behavior_mode}")
         self.vivit = ViViT(args, input_shape=tuple(input_shape))
+        self.vivit.freeze_non_transformer_layers()
         self.output_shape = self.vivit.output_shape
 
     def regularizer(self):
@@ -1041,3 +1121,4 @@ def forward(
             )
         outputs = outputs.to(torch.float32)
         return outputs
+