birder 0.4.1__py3-none-any.whl → 0.4.4__py3-none-any.whl

birder/net/nextvit.py

@@ -355,18 +355,18 @@ class NextViT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
         # Weights initialization
         for m in self.modules():
             if isinstance(m, (nn.BatchNorm2d, nn.GroupNorm, nn.LayerNorm, nn.BatchNorm1d)):
-                nn.init.constant_(m.weight, 1.0)
-                nn.init.constant_(m.bias, 0)
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
 
             elif isinstance(m, nn.Linear):
                 nn.init.normal_(m.weight, std=0.02)
                 if hasattr(m, "bias") and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
+                    nn.init.zeros_(m.bias)
 
             elif isinstance(m, nn.Conv2d):
                 nn.init.normal_(m.weight, std=0.02)
                 if hasattr(m, "bias") and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
+                    nn.init.zeros_(m.bias)
 
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
         x = self.stem(x)

birder/net/resnext.py

@@ -205,8 +205,8 @@ class ResNeXt(DetectorBackbone):
                 nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
 
             elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
 
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
         x = self.stem(x)

birder/net/rope_deit3.py

@@ -223,7 +223,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
                 pt_grid_size=self.pt_grid_size,
             ),
             dim=-1,
-        ).to(self.rope.pos_embed.device)
+        ).to(self.rope.pos_embed.device, dtype=self.rope.pos_embed.dtype)
 
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
         H, W = x.shape[-2:]
@@ -249,7 +249,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
             xs = self.encoder.forward_features(x, rope, out_indices=self.out_indices)
 
         out: dict[str, torch.Tensor] = {}
-        for stage_name, stage_x in zip(self.return_stages, xs):
+        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
             stage_x = stage_x[:, self.num_special_tokens :]
             stage_x = stage_x.permute(0, 2, 1)
             B, C, _ = stage_x.size()

birder/net/rope_flexivit.py

@@ -292,7 +292,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
                 pt_grid_size=self.pt_grid_size,
             ),
             dim=-1,
-        ).to(self.rope.pos_embed.device)
+        ).to(self.rope.pos_embed.device, dtype=self.rope.pos_embed.dtype)
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
@@ -342,7 +342,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             xs = self.encoder.forward_features(x, rope, out_indices=self.out_indices)
 
         out: dict[str, torch.Tensor] = {}
-        for stage_name, stage_x in zip(self.return_stages, xs):
+        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
             stage_x = stage_x[:, self.num_special_tokens :]
             stage_x = stage_x.permute(0, 2, 1)
             B, C, _ = stage_x.size()

birder/net/rope_vit.py

@@ -648,7 +648,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
                 pt_grid_size=self.pt_grid_size,
             ),
             dim=-1,
-        ).to(self.rope.pos_embed.device)
+        ).to(self.rope.pos_embed.device, dtype=self.rope.pos_embed.dtype)
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
@@ -698,7 +698,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             xs = self.encoder.forward_features(x, rope, out_indices=self.out_indices)
 
         out: dict[str, torch.Tensor] = {}
-        for stage_name, stage_x in zip(self.return_stages, xs):
+        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
             stage_x = stage_x[:, self.num_special_tokens :]
             stage_x = stage_x.permute(0, 2, 1)
             B, C, _ = stage_x.size()

birder/net/simple_vit.py

@@ -215,7 +215,7 @@ class Simple_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin):
             xs = self.encoder.forward_features(x, out_indices=self.out_indices)
 
         out: dict[str, torch.Tensor] = {}
-        for stage_name, stage_x in zip(self.return_stages, xs):
+        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
             stage_x = stage_x[:, self.num_special_tokens :]
             stage_x = stage_x.permute(0, 2, 1)
             B, C, _ = stage_x.size()

birder/net/squeezenet.py

@@ -76,7 +76,7 @@ class SqueezeNet(BaseNet):
             if isinstance(m, nn.Conv2d):
                 nn.init.kaiming_uniform_(m.weight)
                 if m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
+                    nn.init.zeros_(m.bias)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/ssl/capi.py

@@ -66,14 +66,20 @@ def sinkhorn_knopp_(M: torch.Tensor, temp: float, n_iterations: int, eps: float
 
 
 class SinkhornQueue(nn.Module):
-    def __init__(self, queue_size: int, position_wise: bool) -> None:
+    def __init__(self, queue_size: int, position_wise: bool, dim: int, seq_len: Optional[int] = None) -> None:
         super().__init__()
         self.queue_size = queue_size
         self.position_wise = position_wise
         self.active = True
-        self.queue = nn.Buffer(torch.empty(0), persistent=False)
-        self.queue_ptr: int = 0
-        self.queue_full: bool = False
+        if self.position_wise is True:
+            assert seq_len is not None, "seq_len is required when position_wise is True"
+
+            self.queue = nn.Buffer(torch.empty(seq_len, queue_size, dim))
+        else:
+            self.queue = nn.Buffer(torch.empty(queue_size, dim))
+
+        self.queue_ptr = nn.Buffer(torch.zeros(1, dtype=torch.long))
+        self.queue_full = nn.Buffer(torch.zeros(1, dtype=torch.bool))
 
     def set_active(self, active: bool) -> None:
         self.active = active
@@ -81,13 +87,13 @@ class SinkhornQueue(nn.Module):
     def get(self) -> Optional[torch.Tensor]:
         if self.active is False:
             return None
-        if self.queue_full is False:
+        if self.queue_full.item() is False:
             return None
 
         return self.queue
 
     @torch.no_grad()  # type: ignore[untyped-decorator]
-    def forward(self, values: torch.Tensor) -> None:  # pylint: disable=too-many-branches
+    def forward(self, values: torch.Tensor) -> None:
         if self.active is False:
             return
         if values.numel() == 0:
@@ -98,21 +104,16 @@ class SinkhornQueue(nn.Module):
             if values.dim() != 3:
                 raise ValueError("SinkhornQueue in position wise mode expects a 3D tensor")
 
-            seq_len = values.size(0)
             batch_size = values.size(1)
-            dim = values.size(2)
-
-            if self.queue.numel() == 0:
-                self.queue = values.new_empty(seq_len, self.queue_size, dim)
 
             values = values.detach()
             if batch_size >= self.queue_size:
                 self.queue.copy_(values[:, -self.queue_size :, :])
-                self.queue_ptr = 0
-                self.queue_full = True
+                self.queue_ptr.zero_()
+                self.queue_full.fill_(True)
                 return
 
-            ptr = self.queue_ptr
+            ptr = self.queue_ptr.item()
             end = ptr + batch_size
             if end <= self.queue_size:
                 self.queue[:, ptr:end, :].copy_(values)
@@ -121,26 +122,23 @@ class SinkhornQueue(nn.Module):
                 self.queue[:, ptr:, :].copy_(values[:, :first, :])
                 self.queue[:, : end - self.queue_size, :].copy_(values[:, first:, :])
 
-            self.queue_ptr = end % self.queue_size
+            self.queue_ptr.fill_(end % self.queue_size)
             if end >= self.queue_size:
-                self.queue_full = True
+                self.queue_full.fill_(True)
 
         else:
             # values shape: (N, dim) - 2D
             if values.dim() != 2:
                 raise ValueError("SinkhornQueue in non-position wise mode expects a 2D tensor")
 
-            if self.queue.numel() == 0:
-                self.queue = values.new_empty(self.queue_size, values.size(1))
-
             values = values.detach()
             if values.size(0) >= self.queue_size:
                 self.queue.copy_(values[-self.queue_size :])
-                self.queue_ptr = 0
-                self.queue_full = True
+                self.queue_ptr.zero_()
+                self.queue_full.fill_(True)
                 return
 
-            ptr = self.queue_ptr
+            ptr = self.queue_ptr.item()
             end = ptr + values.size(0)
             if end <= self.queue_size:
                 self.queue[ptr:end].copy_(values)
@@ -149,9 +147,9 @@ class SinkhornQueue(nn.Module):
                 self.queue[ptr:].copy_(values[:first])
                 self.queue[: end - self.queue_size].copy_(values[first:])
 
-            self.queue_ptr = end % self.queue_size
+            self.queue_ptr.fill_(end % self.queue_size)
             if end >= self.queue_size:
-                self.queue_full = True
+                self.queue_full.fill_(True)
 
 
 class OnlineClustering(nn.Module):
@@ -166,6 +164,7 @@ class OnlineClustering(nn.Module):
         pred_temp: float,
         position_wise_sk: bool = True,
         queue_size: Optional[int] = None,
+        seq_len: Optional[int] = None,
     ):
         super().__init__()
         self.n_sk_iter = n_sk_iter
@@ -176,7 +175,9 @@ class OnlineClustering(nn.Module):
         if queue_size is None:
             self.sinkhorn_queue = None
         else:
-            self.sinkhorn_queue = SinkhornQueue(queue_size, position_wise=position_wise_sk)
+            self.sinkhorn_queue = SinkhornQueue(
+                queue_size, position_wise=position_wise_sk, dim=out_dim, seq_len=seq_len
+            )
 
         # Weight initialization
         nn.init.normal_(self.layer.weight, std=1.0)
@@ -399,6 +400,11 @@ class CAPITeacher(SSLBaseNet):
         sk_mode: str = self.config["sk_mode"]
         queue_size: Optional[int] = self.config.get("queue_size", None)
 
+        queue_seq_len: Optional[int] = None
+        if sk_mode == "position-wise" and queue_size is not None:
+            input_size = (self.size[0] // self.backbone.max_stride, self.size[1] // self.backbone.max_stride)
+            queue_seq_len = input_size[0] * input_size[1]
+
         self.head = OnlineClustering(
             self.backbone.embedding_size,
             num_clusters,
@@ -408,6 +414,7 @@ class CAPITeacher(SSLBaseNet):
             pred_temp=pred_temp,
             position_wise_sk=sk_mode == "position-wise",
             queue_size=queue_size,
+            seq_len=queue_seq_len,
         )
 
     def forward(  # type: ignore[override]  # pylint: disable=arguments-differ

birder/net/ssl/dino_v2.py

@@ -76,13 +76,13 @@ class DINOHead(nn.Module):
 
 
 class SinkhornQueue(nn.Module):
-    def __init__(self, queue_size: int) -> None:
+    def __init__(self, queue_size: int, dim: int) -> None:
         super().__init__()
         self.queue_size = queue_size
         self.active = True
-        self.queue = nn.Buffer(torch.empty(0), persistent=False)
-        self.queue_ptr: int = 0
-        self.queue_full: bool = False
+        self.queue = nn.Buffer(torch.empty(queue_size, dim))
+        self.queue_ptr = nn.Buffer(torch.zeros(1, dtype=torch.long))
+        self.queue_full = nn.Buffer(torch.zeros(1, dtype=torch.bool))
 
     def set_active(self, active: bool) -> None:
         self.active = active
@@ -90,7 +90,7 @@ class SinkhornQueue(nn.Module):
     def get(self) -> Optional[torch.Tensor]:
         if self.active is False:
             return None
-        if self.queue_full is False:
+        if self.queue_full.item() is False:
             return None
 
         return self.queue
@@ -104,17 +104,14 @@ class SinkhornQueue(nn.Module):
         if values.dim() != 2:
             raise ValueError("SinkhornQueue expects a 2D tensor")
 
-        if self.queue.numel() == 0:
-            self.queue = values.new_empty(self.queue_size, values.size(1))
-
         values = values.detach()
         if values.size(0) >= self.queue_size:
             self.queue.copy_(values[-self.queue_size :])
-            self.queue_ptr = 0
-            self.queue_full = True
+            self.queue_ptr.zero_()
+            self.queue_full.fill_(True)
             return
 
-        ptr = self.queue_ptr
+        ptr = self.queue_ptr.item()
         end = ptr + values.size(0)
         if end <= self.queue_size:
             self.queue[ptr:end].copy_(values)
@@ -123,9 +120,9 @@ class SinkhornQueue(nn.Module):
             self.queue[ptr:].copy_(values[:first])
             self.queue[: end - self.queue_size].copy_(values[first:])
 
-        self.queue_ptr = end % self.queue_size
+        self.queue_ptr.fill_(end % self.queue_size)
         if end >= self.queue_size:
-            self.queue_full = True
+            self.queue_full.fill_(True)
 
 
 class DINOLoss(nn.Module):
@@ -139,7 +136,7 @@ class DINOLoss(nn.Module):
         if queue_size is None:
             self.sinkhorn_queue = None
         else:
-            self.sinkhorn_queue = SinkhornQueue(queue_size)
+            self.sinkhorn_queue = SinkhornQueue(queue_size, dim=out_dim)
 
         self.updated = True
         self.reduce_handle: Any = None
@@ -267,7 +264,7 @@ class iBOTPatchLoss(nn.Module):
         if queue_size is None:
             self.sinkhorn_queue = None
         else:
-            self.sinkhorn_queue = SinkhornQueue(queue_size)
+            self.sinkhorn_queue = SinkhornQueue(queue_size, dim=patch_out_dim)
 
         self.updated = True
         self.reduce_handle: Any = None

birder/net/ssl/franca.py

@@ -124,13 +124,13 @@ class DINOHeadMRL(nn.Module):
 
 
 class SinkhornQueue(nn.Module):
-    def __init__(self, queue_size: int) -> None:
+    def __init__(self, queue_size: int, dim: int) -> None:
         super().__init__()
         self.queue_size = queue_size
         self.active = True
-        self.queue = nn.Buffer(torch.empty(0), persistent=False)
-        self.queue_ptr: int = 0
-        self.queue_full: bool = False
+        self.queue = nn.Buffer(torch.empty(queue_size, dim))
+        self.queue_ptr = nn.Buffer(torch.zeros(1, dtype=torch.long))
+        self.queue_full = nn.Buffer(torch.zeros(1, dtype=torch.bool))
 
     def set_active(self, active: bool) -> None:
         self.active = active
@@ -138,7 +138,7 @@ class SinkhornQueue(nn.Module):
     def get(self) -> Optional[torch.Tensor]:
         if self.active is False:
             return None
-        if self.queue_full is False:
+        if self.queue_full.item() is False:
             return None
 
         return self.queue
@@ -152,17 +152,14 @@ class SinkhornQueue(nn.Module):
         if values.dim() != 2:
             raise ValueError("SinkhornQueue expects a 2D tensor")
 
-        if self.queue.numel() == 0:
-            self.queue = values.new_empty(self.queue_size, values.size(1))
-
         values = values.detach()
         if values.size(0) >= self.queue_size:
             self.queue.copy_(values[-self.queue_size :])
-            self.queue_ptr = 0
-            self.queue_full = True
+            self.queue_ptr.zero_()
+            self.queue_full.fill_(True)
             return
 
-        ptr = self.queue_ptr
+        ptr = self.queue_ptr.item()
         end = ptr + values.size(0)
         if end <= self.queue_size:
             self.queue[ptr:end].copy_(values)
@@ -171,13 +168,15 @@ class SinkhornQueue(nn.Module):
             self.queue[ptr:].copy_(values[:first])
             self.queue[: end - self.queue_size].copy_(values[first:])
 
-        self.queue_ptr = end % self.queue_size
+        self.queue_ptr.fill_(end % self.queue_size)
         if end >= self.queue_size:
-            self.queue_full = True
+            self.queue_full.fill_(True)
 
 
 class DINOLossMRL(nn.Module):
-    def __init__(self, student_temp: float, nesting_levels: int, queue_size: Optional[int] = None) -> None:
+    def __init__(
+        self, student_temp: float, nesting_levels: int, queue_size: Optional[int] = None, out_dim: Optional[int] = None
+    ) -> None:
         super().__init__()
         self.student_temp = student_temp
         self.queue_active = True
@@ -185,9 +184,12 @@ class DINOLossMRL(nn.Module):
         if queue_size is None:
             self.sinkhorn_queue = None
         else:
+            assert out_dim is not None, "out_dim is required when queue_size is set"
+
+            queue_dims = _get_nesting_list(out_dim, nesting_levels)
             self.sinkhorn_queue = nn.ModuleList()
-            for _ in range(nesting_levels):
-                queue = SinkhornQueue(queue_size)
+            for dim in queue_dims:
+                queue = SinkhornQueue(queue_size, dim)
                 queue.set_active(self.queue_active)
                 self.sinkhorn_queue.append(queue)
 
@@ -300,7 +302,9 @@ class DINOLossMRL(nn.Module):
 
 # pylint: disable=invalid-name
 class iBOTPatchLossMRL(nn.Module):
-    def __init__(self, student_temp: float, nesting_levels: int, queue_size: Optional[int] = None) -> None:
+    def __init__(
+        self, student_temp: float, nesting_levels: int, queue_size: Optional[int] = None, out_dim: Optional[int] = None
+    ) -> None:
         super().__init__()
         self.student_temp = student_temp
         self.queue_active = True
@@ -308,9 +312,12 @@ class iBOTPatchLossMRL(nn.Module):
         if queue_size is None:
             self.sinkhorn_queue = None
         else:
+            assert out_dim is not None, "out_dim is required when queue_size is set"
+
+            queue_dims = _get_nesting_list(out_dim, nesting_levels)
             self.sinkhorn_queue = nn.ModuleList()
-            for _ in range(nesting_levels):
-                queue = SinkhornQueue(queue_size)
+            for dim in queue_dims:
+                queue = SinkhornQueue(queue_size, dim)
                 queue.set_active(self.queue_active)
                 self.sinkhorn_queue.append(queue)
 

birder/net/van.py

@@ -206,11 +206,11 @@ class VAN(DetectorBackbone):
             if isinstance(m, nn.Linear):
                 nn.init.trunc_normal_(m.weight, std=0.02)
                 if m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
+                    nn.init.zeros_(m.bias)
 
             elif isinstance(m, nn.LayerNorm):
+                nn.init.ones_(m.weight)
                 nn.init.zeros_(m.bias)
-                nn.init.constant_(m.weight, 1.0)
 
             elif isinstance(m, nn.Conv2d):
                 fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels

birder/net/vit.py

@@ -572,7 +572,7 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
             xs = self.encoder.forward_features(x, out_indices=self.out_indices)
 
         out: dict[str, torch.Tensor] = {}
-        for stage_name, stage_x in zip(self.return_stages, xs):
+        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
             stage_x = stage_x[:, self.num_special_tokens :]
             stage_x = stage_x.permute(0, 2, 1)
             B, C, _ = stage_x.size()
@@ -802,6 +802,24 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
 # Register model configs (side effects)
 register_vit_configs(ViT)
 
+registry.register_weights(  # BioCLIP v1: https://arxiv.org/abs/2311.18803
+    "vit_b16_pn_bioclip-v1",
+    {
+        "url": "https://huggingface.co/birder-project/vit_b16_pn_bioclip-v1/resolve/main",
+        "description": (
+            "ViT b16 image encoder pre-trained by Imageomics using CLIP on the TreeOfLife-10M dataset. "
+            "This model has not been fine-tuned for a specific classification task"
+        ),
+        "resolution": (224, 224),
+        "formats": {
+            "pt": {
+                "file_size": 328.9,
+                "sha256": "9b2e5598f233657932eeb77e027cd4c4d683bf75515768fe6971cab6ec10bf15",
+            },
+        },
+        "net": {"network": "vit_b16_pn", "tag": "bioclip-v1"},
+    },
+)
 registry.register_weights(
     "vit_l16_mim_200",
     {
@@ -849,8 +867,8 @@ registry.register_weights(  # BioCLIP v2: https://arxiv.org/abs/2505.23883
         "resolution": (224, 224),
         "formats": {
             "pt": {
-                "file_size": 1156.6,
-                "sha256": "6cd7bd6993762590891fe2b41db1649cde5a0c4de5a7f341672f8856ed529d07",
+                "file_size": 1159.7,
+                "sha256": "301a325579dafdfa2ea13b0cbaf8129211ecd1429c29afa20d1c2eaaa91d8b0d",
             },
         },
         "net": {"network": "vit_l14_pn", "tag": "bioclip-v2"},

birder/net/vit_parallel.py

@@ -370,7 +370,7 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             xs = self.encoder.forward_features(x, out_indices=self.out_indices)
 
         out: dict[str, torch.Tensor] = {}
-        for stage_name, stage_x in zip(self.return_stages, xs):
+        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
             stage_x = stage_x[:, self.num_special_tokens :]
             stage_x = stage_x.permute(0, 2, 1)
             B, C, _ = stage_x.size()

birder/net/vit_sam.py

@@ -29,7 +29,9 @@ from birder.net._vit_configs import BASE
 from birder.net._vit_configs import HUGE
 from birder.net._vit_configs import LARGE
 from birder.net._vit_configs import MEDIUM
+from birder.net._vit_configs import SMALL
 from birder.net.base import DetectorBackbone
+from birder.net.base import normalize_out_indices
 from birder.net.vit import EncoderBlock as MAEDecoderBlock
 
 
@@ -72,7 +74,7 @@ def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor
 
     # Interpolate rel pos if needed
     if rel_pos.shape[0] != max_rel_dist:
-        # Adjust size is a one off interpolation, should prevent us from getting here
+        # Only reached in dynamic-size mode (rel-pos table resized on the fly)
         rel_pos_resized = F.interpolate(
             rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), size=max_rel_dist, mode="linear"
         )
@@ -242,6 +244,7 @@ class EncoderBlock(nn.Module):
 class ViT_SAM(DetectorBackbone):
     block_group_regex = r"body\.(\d+)"
 
+    # pylint: disable=too-many-locals
     def __init__(
         self,
         input_channels: int,
@@ -266,6 +269,7 @@ class ViT_SAM(DetectorBackbone):
         window_size: int = self.config["window_size"]
         global_attn_indexes: list[int] = self.config["global_attn_indexes"]
         neck_channels: Optional[int] = self.config.get("neck_channels", None)
+        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
         drop_path_rate: float = self.config["drop_path_rate"]
 
         if norm_layer_type == "LayerNorm":
@@ -292,6 +296,7 @@ class ViT_SAM(DetectorBackbone):
         self.hidden_dim = hidden_dim
         self.global_attn_indexes = global_attn_indexes
         self.num_special_tokens = 0
+        self.out_indices = normalize_out_indices(out_indices, num_layers)
         dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]  # Stochastic depth decay rule
 
         self.patch_embed = PatchEmbed(
@@ -356,8 +361,10 @@ class ViT_SAM(DetectorBackbone):
             nn.Flatten(1),
         )
 
-        self.return_stages = ["neck"]  # Actually meaningless, but for completeness
-        self.return_channels = [neck_channels]
+        num_return_stages = len(self.out_indices) if self.out_indices is not None else 1
+        self.return_stages = [f"stage{stage_idx + 1}" for stage_idx in range(num_return_stages)]
+        self.return_channels = [hidden_dim] * num_return_stages
+        self.return_channels[-1] = neck_channels
         self.embedding_size = neck_channels
         self.classifier = self.create_classifier()
 
@@ -372,13 +379,54 @@ class ViT_SAM(DetectorBackbone):
             activation_layer=nn.GELU,
         )
 
+    def _get_pos_embed(self, H: int, W: int) -> torch.Tensor:
+        if self.dynamic_size is False:
+            return self.pos_embedding
+
+        if H == self.size[0] and W == self.size[1]:
+            return self.pos_embedding
+
+        base_h = H // self.patch_size
+        base_w = W // self.patch_size
+        orig_dtype = self.pos_embedding.dtype
+        pos_embedding = self.pos_embedding.float()
+        pos_embedding = pos_embedding.permute(0, 3, 1, 2)
+        pos_embedding = F.interpolate(pos_embedding, size=(base_h, base_w), mode="bicubic", antialias=True)
+        pos_embedding = pos_embedding.permute(0, 2, 3, 1)
+
+        return pos_embedding.to(orig_dtype)
+
+    def set_causal_attention(self, is_causal: bool = True) -> None:
+        for b in self.body:
+            b.set_causal_attention(is_causal)
+
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
+        H, W = x.shape[-2:]
         x = self.patch_embed(x)
-        x = x + self.pos_embedding
+        x = x + self._get_pos_embed(H, W)
 
-        x = self.body(x)
-        x = self.neck(x.permute(0, 3, 1, 2))
-        return {self.return_stages[0]: x}
+        if self.out_indices is None:
+            x = self.body(x)
+            x = self.neck(x.permute(0, 3, 1, 2))
+            return {self.return_stages[0]: x}
+
+        out_indices_set = set(self.out_indices)
+        last_out_idx = max(out_indices_set)
+        out: dict[str, torch.Tensor] = {}
+        stage_idx = 0
+        for idx, blk in enumerate(self.body):
+            x = blk(x)
+            if idx not in out_indices_set:
+                continue
+
+            stage_x = x.permute(0, 3, 1, 2)
+            if idx == last_out_idx:
+                stage_x = self.neck(stage_x)
+
+            out[self.return_stages[stage_idx]] = stage_x
+            stage_idx += 1
+
+        return out
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.patch_embed.parameters():
@@ -393,13 +441,10 @@ class ViT_SAM(DetectorBackbone):
             for param in module.parameters():
                 param.requires_grad_(False)
 
-    def set_causal_attention(self, is_causal: bool = True) -> None:
-        for b in self.body:
-            b.set_causal_attention(is_causal)
-
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        H, W = x.shape[-2:]
         x = self.patch_embed(x)
-        x = x + self.pos_embedding
+        x = x + self._get_pos_embed(H, W)
 
         x = self.body(x)
         x = self.neck(x.permute(0, 3, 1, 2))
@@ -410,9 +455,6 @@ class ViT_SAM(DetectorBackbone):
         x = self.forward_features(x)
         return self.features(x)
 
-    def set_dynamic_size(self, dynamic_size: bool = True) -> None:
-        assert dynamic_size is False, "Dynamic size not supported for this network"
-
     def adjust_size(self, new_size: tuple[int, int]) -> None:
         if new_size == self.size:
             return
@@ -530,6 +572,11 @@ class ViT_SAM(DetectorBackbone):
 
 
 # ViTDet (no neck)
+registry.register_model_config(
+    "vit_det_s16",
+    ViT_SAM,
+    config={"patch_size": 16, **SMALL, "window_size": 14, "global_attn_indexes": [2, 5, 8, 11]},
+)
 registry.register_model_config(
     "vit_det_m16_rms",
     ViT_SAM,
@@ -541,7 +588,6 @@ registry.register_model_config(
         "global_attn_indexes": [2, 5, 8, 11],
     },
 )
-
 registry.register_model_config(
     "vit_det_b16",
     ViT_SAM,