Encoder2 is now default

Author: Tensor46

examples/finetune_minivit_imagenet.yaml

@@ -32,7 +32,7 @@ backbone:
   first_stride: 2
   last_pool: half
   last_stride: 4
-  use_sincos_pos_token: True
+  pos_token: "sincos"
 
   mode: full
   skip_freeze_prefixes:

examples/mae_minivit_imagenet.yaml

@@ -23,7 +23,7 @@ mae_encoder:
   first_stride: 2
   last_pool: half
   last_stride: 4
-  use_sincos_pos_token: False
+  pos_token: "sincos"
 
 mae_decoder:
   name: MAEDecoder
@@ -37,6 +37,7 @@ mae_decoder:
   mlp_dim: 2048
   num_layers: 4
   num_heads: 16
+  pos_token: "sincos"
 
 loss:
   name: MSELoss

src/sensa/models/mae_decoder.py

@@ -1,9 +1,10 @@
 from collections.abc import Callable
 from functools import partial
+from typing import Literal
 
 import torch
 
-from sensa.layers.encoder import Encoder
+from sensa.layers.encoder import Encoder2
 from sensa.models.base import BaseModel
 from sensa.models.registry import register_model
 
@@ -29,8 +30,12 @@ class MAEDecoder(BaseModel):
             Number of transformer layers in the decoder.
         num_heads (int):
             Number of attention heads in each decoder layer.
+        act_layer (Callable[..., torch.nn.Module]):
+            Activation layer for the decoder. Defaults to torch.nn.GELU.
         norm_layer (Callable[..., torch.nn.Module], optional):
             Constructor for the normalization layer. Defaults to `partial(torch.nn.LayerNorm, eps=1e-6)`.
+        pos_token (Literal["learned", "sincos", "rope"]):
+            Positional token type. Defaults to "sincos".
     """
 
     def __init__(
@@ -43,7 +48,9 @@ def __init__(
         mlp_dim: int,
         num_layers: int,
         num_heads: int,
+        act_layer: Callable[..., torch.nn.Module] = torch.nn.GELU,
         norm_layer: Callable[..., torch.nn.Module] | None = None,
+        pos_token: Literal["learned", "sincos", "rope"] = "sincos",
     ):
         super().__init__()
         self.image_size = image_size
@@ -55,18 +62,19 @@ def __init__(
             torch.nn.Identity() if encoder_dim == decoder_dim else torch.nn.Linear(encoder_dim, decoder_dim)
         )
         # build the decoder transformer
-        self.decoder = Encoder(
+        self.decoder = Encoder2(
             size=self.stem_size,
             extra_tokens=0,
             num_layers=num_layers,
             num_heads=num_heads,
             hidden_dim=decoder_dim,
             mlp_dim=mlp_dim,
             dropout=0.0,
-            attention_dropout=0.0,
+            act_layer=act_layer,
             norm_layer=partial(torch.nn.LayerNorm, eps=1e-6) if norm_layer is None else norm_layer,
+            pos_token=pos_token,
         )
-        self.decoder.use_sincos_pos_token(extra_tokens=0, size=self.stem_size)
+        # self.decoder.use_sincos_pos_token(extra_tokens=0, size=self.stem_size)
         # projection head to map decoder outputs back to patch pixels
         self.predict = torch.nn.Linear(decoder_dim, patch_size * patch_size * channels)
 

src/sensa/models/vit.py

@@ -11,7 +11,7 @@
 
 from sensa.layers import mask_utils
 from sensa.layers.dyt import DyT
-from sensa.layers.encoder import Encoder
+from sensa.layers.encoder import Encoder2
 from sensa.layers.last_pool import LastPool
 from sensa.models.base import BaseModel
 from sensa.models.registry import register_model
@@ -157,10 +157,12 @@ class VIT(BaseModel):
                 Defaults to "token".
         last_stride (int, optional):
             Stride for the stem's final downsampling block. Defaults to 4.
+        act_layer (Callable[..., torch.nn.Module]):
+            Activation layer for encoder. Defaults to torch.nn.GELU.
         norm_layer (Callable[..., nn.Module] | str, optional):
             Normalization layer for encoder. Defaults to LayerNorm(eps=1e-6).
-        use_sincos_pos_token (bool, optional):
-            Whether to use fixed sinusoidal positional embeddings. Defaults to False.
+        pos_token (Literal["learned", "sincos", "rope"]):
+            Positional token type. Defaults to "learned".
     """
 
     def __init__(
@@ -177,8 +179,9 @@ def __init__(
         first_stride: int = 2,
         last_pool: Literal["avg", "full", "half", "token", None] = "token",
         last_stride: int = 4,
+        act_layer: Callable[..., torch.nn.Module] = torch.nn.GELU,
         norm_layer: Callable[..., torch.nn.Module] | str | None = None,
-        use_sincos_pos_token: bool = False,
+        pos_token: Literal["learned", "sincos", "rope"] = "learned",
     ):
         super().__init__()
         self.image_size = torch.nn.modules.utils._pair(image_size)
@@ -216,19 +219,20 @@ def __init__(
             self.class_token = torch.nn.Parameter(torch.zeros(1, 1, hidden_dim))
             extra_tokens += 1
 
-        self.encoder = Encoder(
+        self.encoder = Encoder2(
             size=self.stem_size,
             extra_tokens=extra_tokens,
             num_layers=num_layers,
             num_heads=num_heads,
             hidden_dim=hidden_dim,
             mlp_dim=mlp_dim,
             dropout=0.0,
-            attention_dropout=0.0,
+            act_layer=act_layer,
             norm_layer=norm_layer,
+            pos_token=pos_token,
         )
-        if use_sincos_pos_token:
-            self.encoder.use_sincos_pos_token(extra_tokens=int(last_pool == "token"), size=self.stem_size)
+        # if pos_token == "sincos":
+        #     self.encoder.use_sincos_pos_token(extra_tokens=int(last_pool == "token"), size=self.stem_size)
         self.seq_length = self.encoder.seq_length
 
         if self.mask_ratio > 0:
@@ -310,7 +314,11 @@ def param_groups(self) -> list[dict[str, Any]]:
         for i in range(0, len(self.encoder.layers), 2):
             groups += self._param_groups(self.encoder.layers[slice(i, i + 2)])
         self._param_groups(self.encoder.ln, groups=groups[-2:])
-        if isinstance(self.encoder.pos_token, torch.nn.Parameter) and self.encoder.pos_token.requires_grad:
+        if (
+            hasattr(self.encoder, "pos_token")
+            and isinstance(self.encoder.pos_token, torch.nn.Parameter)
+            and self.encoder.pos_token.requires_grad
+        ):
             groups[-1]["params"].append(self.encoder.pos_token)
         if (
             hasattr(self, "class_token")

tests/samples/mae_vit.yaml

@@ -24,7 +24,7 @@ mae_encoder:
   last_pool: half
   last_stride: 4
   norm_layer: dyt
-  use_sincos_pos_token: True
+  pos_token: "sincos"
 
 mae_decoder:
   name: MAEDecoder

tests/test_models.py

@@ -47,7 +47,7 @@ def test_vit_features():
         num_classes=None,
         in_channels=3,
         last_pool=None,
-        use_sincos_pos_token=True,
+        pos_token="sincos",
     )
     output = model(torch.randn(1, 3, 128, 128))
     assert output.shape[-1] == model.hidden_dim, f"output shape must be {output.shape}"
@@ -67,7 +67,7 @@ def test_vit_features_with_pool():
         num_classes=None,
         in_channels=3,
         last_pool="half",
-        use_sincos_pos_token=True,
+        pos_token="rope",
     )
     output = model(torch.randn(1, 3, 128, 128))
     size = model.hidden_dim * (model.stem_size[0] // 2) * (model.stem_size[1] // 2)