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

birder/net/repvgg.py

@@ -56,7 +56,6 @@ class RepVggBlock(nn.Module):
                 stride=(stride, stride),
                 padding=(padding, padding),
                 groups=groups,
-                bias=True,
             )
         else:
             self.reparam_conv = None
@@ -113,7 +112,7 @@ class RepVggBlock(nn.Module):
         if self.reparameterized is True:
             return
 
-        (kernel, bias) = self._get_kernel_bias()
+        kernel, bias = self._get_kernel_bias()
         self.reparam_conv = nn.Conv2d(
             in_channels=self.conv_kxk.conv.in_channels,
             out_channels=self.conv_kxk.conv.out_channels,
@@ -122,7 +121,6 @@ class RepVggBlock(nn.Module):
             padding=self.conv_kxk.conv.padding,
             dilation=self.conv_kxk.conv.dilation,
             groups=self.conv_kxk.conv.groups,
-            bias=True,
         )
         self.reparam_conv.weight.data = kernel
         self.reparam_conv.bias.data = bias
@@ -151,10 +149,10 @@ class RepVggBlock(nn.Module):
         kernel_identity = 0
         bias_identity = 0
         if self.rbr_identity is not None:
-            (kernel_identity, bias_identity) = self._fuse_bn_tensor(self.rbr_identity)
+            kernel_identity, bias_identity = self._fuse_bn_tensor(self.rbr_identity)
 
         # Get weights and bias of conv branches
-        (kernel_conv, bias_conv) = self._fuse_bn_tensor(self.conv_kxk)
+        kernel_conv, bias_conv = self._fuse_bn_tensor(self.conv_kxk)
 
         kernel_final = kernel_conv + kernel_1x1 + kernel_identity
         bias_final = bias_conv + bias_1x1 + bias_identity

birder/net/repvit.py

@@ -60,7 +60,7 @@ class RepConvBN(nn.Sequential):
         if self.reparameterized is True:
             return
 
-        (c, bn) = self._modules.values()
+        c, bn = self._modules.values()
         w = bn.weight / (bn.running_var + bn.eps) ** 0.5
         w = c.weight * w[:, None, None, None]
         b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5
@@ -101,7 +101,7 @@ class RepNormLinear(nn.Sequential):
         if self.reparameterized is True:
             return
 
-        (bn, li) = self._modules.values()
+        bn, li = self._modules.values()
         w = bn.weight / (bn.running_var + bn.eps) ** 0.5
         b = bn.bias - self.bn.running_mean * self.bn.weight / (bn.running_var + bn.eps) ** 0.5
         w = li.weight * w[None, :]

birder/net/resnest.py

@@ -85,7 +85,7 @@ class SplitAttn(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.conv(x)
 
-        (B, RC, H, W) = x.size()  # pylint: disable=invalid-name
+        B, RC, H, W = x.size()  # pylint: disable=invalid-name
         if self.radix > 1:
             x = x.reshape((B, self.radix, RC // self.radix, H, W))
             x_gap = x.sum(dim=1)

birder/net/rope_deit3.py

@@ -34,6 +34,7 @@ from birder.net.base import MaskedTokenRetentionMixin
 from birder.net.base import PreTrainEncoder
 from birder.net.base import TokenOmissionResultType
 from birder.net.base import TokenRetentionResultType
+from birder.net.base import normalize_out_indices
 from birder.net.rope_vit import Encoder
 from birder.net.rope_vit import MAEDecoderBlock
 from birder.net.rope_vit import RoPE
@@ -46,6 +47,7 @@ from birder.net.vit import adjust_position_embedding
 class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTokenRetentionMixin):
     block_group_regex = r"encoder\.block\.(\d+)"
 
+    # pylint: disable=too-many-locals
     def __init__(
         self,
         input_channels: int,
@@ -68,6 +70,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
         mlp_dim: int = self.config["mlp_dim"]
         layer_scale_init_value: Optional[float] = self.config.get("layer_scale_init_value", 1e-5)
         num_reg_tokens: int = self.config.get("num_reg_tokens", 0)
+        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
         rope_rot_type: Literal["standard", "interleaved"] = self.config.get("rope_rot_type", "standard")
         rope_grid_indexing: Literal["ij", "xy"] = self.config.get("rope_grid_indexing", "ij")
         rope_grid_offset: int = self.config.get("rope_grid_offset", 0)
@@ -86,6 +89,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
         self.num_reg_tokens = num_reg_tokens
         self.num_special_tokens = 1 + self.num_reg_tokens
         self.pos_embed_special_tokens = pos_embed_special_tokens
+        self.out_indices = normalize_out_indices(out_indices, num_layers)
         self.rope_rot_type = rope_rot_type
         self.rope_grid_indexing = rope_grid_indexing
         self.rope_grid_offset = rope_grid_offset
@@ -105,7 +109,6 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
             kernel_size=(patch_size, patch_size),
             stride=(patch_size, patch_size),
             padding=(0, 0),
-            bias=True,
         )
         self.patch_embed = PatchEmbed()
 
@@ -153,8 +156,9 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
         )
         self.norm = nn.LayerNorm(hidden_dim, eps=1e-6)
 
-        self.return_stages = ["neck"]  # Actually meaningless, but for completeness
-        self.return_channels = [hidden_dim]
+        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.embedding_size = hidden_dim
         self.classifier = self.create_classifier()
 
@@ -222,7 +226,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
         ).to(self.rope.pos_embed.device)
 
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
         x = self.conv_proj(x)
         x = self.patch_embed(x)
 
@@ -238,15 +242,21 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
             x = x + self._get_pos_embed(H, W)
             x = torch.concat([batch_special_tokens, x], dim=1)
 
-        x = self.encoder(x, self._get_rope_embed(H, W))
-        x = self.norm(x)
+        rope = self._get_rope_embed(H, W)
+        if self.out_indices is None:
+            xs = [self.encoder(x, rope)]
+        else:
+            xs = self.encoder.forward_features(x, rope, out_indices=self.out_indices)
 
-        x = x[:, self.num_special_tokens :]
-        x = x.permute(0, 2, 1)
-        (B, C, _) = x.size()
-        x = x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+        out: dict[str, torch.Tensor] = {}
+        for stage_name, stage_x in zip(self.return_stages, xs):
+            stage_x = stage_x[:, self.num_special_tokens :]
+            stage_x = stage_x.permute(0, 2, 1)
+            B, C, _ = stage_x.size()
+            stage_x = stage_x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+            out[stage_name] = stage_x
 
-        return {self.return_stages[0]: x}
+        return out
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
@@ -261,6 +271,10 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
             for param in module.parameters():
                 param.requires_grad_(False)
 
+    def transform_to_backbone(self) -> None:
+        super().transform_to_backbone()
+        self.norm = nn.Identity()
+
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)
 
@@ -271,7 +285,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
         return_all_features: bool = False,
         return_keys: Literal["all", "tokens", "embedding"] = "tokens",
     ) -> TokenOmissionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)
@@ -340,7 +354,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
         mask_token: Optional[torch.Tensor] = None,
         return_keys: Literal["all", "features", "embedding"] = "features",
     ) -> TokenRetentionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         x = self.conv_proj(x)
         x = mask_tensor(x, mask, mask_token=mask_token, patch_factor=self.max_stride // self.stem_stride)
@@ -370,7 +384,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
         if return_keys in ("all", "features"):
             features = x[:, self.num_special_tokens :]
             features = features.permute(0, 2, 1)
-            (B, C, _) = features.size()
+            B, C, _ = features.size()
             features = features.reshape(B, C, H // self.patch_size, W // self.patch_size)
             result["features"] = features
 
@@ -380,7 +394,7 @@ class RoPE_DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Ma
         return result
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)

birder/net/rope_flexivit.py

@@ -29,6 +29,7 @@ from birder.net.base import MaskedTokenRetentionMixin
 from birder.net.base import PreTrainEncoder
 from birder.net.base import TokenOmissionResultType
 from birder.net.base import TokenRetentionResultType
+from birder.net.base import normalize_out_indices
 from birder.net.flexivit import flex_proj
 from birder.net.flexivit import get_patch_sizes
 from birder.net.flexivit import interpolate_proj
@@ -82,6 +83,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         norm_layer_eps: float = self.config.get("norm_layer_eps", 1e-6)
         mlp_layer_type: str = self.config.get("mlp_layer_type", "FFN")
         act_layer_type: Optional[str] = self.config.get("act_layer_type", None)  # Default according to mlp type
+        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
         rope_rot_type: Literal["standard", "interleaved"] = self.config.get("rope_rot_type", "standard")
         rope_grid_indexing: Literal["ij", "xy"] = self.config.get("rope_grid_indexing", "ij")
         rope_grid_offset: int = self.config.get("rope_grid_offset", 0)
@@ -125,6 +127,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         self.norm_layer_eps = norm_layer_eps
         self.mlp_layer = mlp_layer
         self.act_layer = act_layer
+        self.out_indices = normalize_out_indices(out_indices, num_layers)
         self.rope_rot_type = rope_rot_type
         self.rope_grid_indexing = rope_grid_indexing
         self.rope_grid_offset = rope_grid_offset
@@ -145,7 +148,6 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             kernel_size=(patch_size, patch_size),
             stride=(patch_size, patch_size),
             padding=(0, 0),
-            bias=True,
         )
         self.patch_embed = PatchEmbed()
 
@@ -218,8 +220,9 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
 
             self.attn_pool = MultiHeadAttentionPool(hidden_dim, attn_pool_num_heads, mlp_dim, qkv_bias=True)
 
-        self.return_stages = ["neck"]  # Actually meaningless, just for completeness
-        self.return_channels = [hidden_dim]
+        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.embedding_size = hidden_dim
         self.classifier = self.create_classifier()
 
@@ -307,8 +310,12 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)
 
+    def transform_to_backbone(self) -> None:
+        super().transform_to_backbone()
+        self.norm = nn.Identity()
+
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
         x = self.conv_proj(x)
         x = self.patch_embed(x)
 
@@ -328,15 +335,21 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         if self.pos_embed_special_tokens is True:
             x = x + self._get_pos_embed(H, W)
 
-        x = self.encoder(x, self._get_rope_embed(H, W))
-        x = self.norm(x)
+        rope = self._get_rope_embed(H, W)
+        if self.out_indices is None:
+            xs = [self.encoder(x, rope)]
+        else:
+            xs = self.encoder.forward_features(x, rope, out_indices=self.out_indices)
 
-        x = x[:, self.num_special_tokens :]
-        x = x.permute(0, 2, 1)
-        (B, C, _) = x.size()
-        x = x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+        out: dict[str, torch.Tensor] = {}
+        for stage_name, stage_x in zip(self.return_stages, xs):
+            stage_x = stage_x[:, self.num_special_tokens :]
+            stage_x = stage_x.permute(0, 2, 1)
+            B, C, _ = stage_x.size()
+            stage_x = stage_x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+            out[stage_name] = stage_x
 
-        return {self.return_stages[0]: x}
+        return out
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
@@ -359,7 +372,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         return_all_features: bool = False,
         return_keys: Literal["all", "tokens", "embedding"] = "tokens",
     ) -> TokenOmissionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)
@@ -439,7 +452,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         mask_token: Optional[torch.Tensor] = None,
         return_keys: Literal["all", "features", "embedding"] = "features",
     ) -> TokenRetentionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         x = self.conv_proj(x)
         x = mask_tensor(x, mask, mask_token=mask_token, patch_factor=self.max_stride // self.stem_stride)
@@ -470,7 +483,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         if return_keys in ("all", "features"):
             features = x[:, self.num_special_tokens :]
             features = features.permute(0, 2, 1)
-            (B, C, _) = features.size()
+            B, C, _ = features.size()
             features = features.reshape(B, C, H // self.patch_size, W // self.patch_size)
             result["features"] = features
 
@@ -490,7 +503,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         return result
 
     def forward_features(self, x: torch.Tensor, patch_size: Optional[int] = None) -> torch.Tensor:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = flex_proj(x, self.conv_proj.weight, self.conv_proj.bias, patch_size)

birder/net/rope_vit.py

@@ -38,6 +38,7 @@ from birder.net.base import MaskedTokenRetentionMixin
 from birder.net.base import PreTrainEncoder
 from birder.net.base import TokenOmissionResultType
 from birder.net.base import TokenRetentionResultType
+from birder.net.base import normalize_out_indices
 from birder.net.vit import PatchEmbed
 from birder.net.vit import adjust_position_embedding
 
@@ -76,7 +77,7 @@ def build_rotary_pos_embed(
 
 def rotate_half(x: torch.Tensor) -> torch.Tensor:
     # Taken from: https://github.com/facebookresearch/capi/blob/main/model.py
-    (x1, x2) = x.chunk(2, dim=-1)
+    x1, x2 = x.chunk(2, dim=-1)
     return torch.concat((-x2, x1), dim=-1)
 
 
@@ -85,7 +86,7 @@ def rotate_half_interleaved(x: torch.Tensor) -> torch.Tensor:
 
 
 def apply_rotary_pos_embed(x: torch.Tensor, embed: torch.Tensor) -> torch.Tensor:
-    (sin_emb, cos_emb) = embed.tensor_split(2, dim=-1)
+    sin_emb, cos_emb = embed.tensor_split(2, dim=-1)
     if cos_emb.ndim == 3:
         return x * cos_emb.unsqueeze(1).expand_as(x) + rotate_half(x) * sin_emb.unsqueeze(1).expand_as(x)
 
@@ -93,7 +94,7 @@ def apply_rotary_pos_embed(x: torch.Tensor, embed: torch.Tensor) -> torch.Tensor
 
 
 def apply_interleaved_rotary_pos_embed(x: torch.Tensor, embed: torch.Tensor) -> torch.Tensor:
-    (sin_emb, cos_emb) = embed.tensor_split(2, dim=-1)
+    sin_emb, cos_emb = embed.tensor_split(2, dim=-1)
     if cos_emb.ndim == 3:
         return x * cos_emb.unsqueeze(1).expand_as(x) + rotate_half_interleaved(x) * sin_emb.unsqueeze(1).expand_as(x)
 
@@ -128,7 +129,7 @@ class RoPE(nn.Module):
         else:
             raise ValueError(f"Unknown rope_rot_type, got '{rope_rot_type}'")
 
-        (sin_emb, cos_emb) = build_rotary_pos_embed(
+        sin_emb, cos_emb = build_rotary_pos_embed(
             dim,
             temperature,
             grid_size=grid_size,
@@ -185,9 +186,9 @@ class RoPEAttention(nn.Module):
         self.proj_drop = nn.Dropout(proj_drop)
 
     def forward(self, x: torch.Tensor, rope: torch.Tensor) -> torch.Tensor:
-        (B, N, C) = x.size()
+        B, N, C = x.size()
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
-        (q, k, v) = qkv.unbind(0)
+        q, k, v = qkv.unbind(0)
         q = self.q_norm(q)
         k = self.k_norm(k)
 
@@ -326,13 +327,17 @@ class Encoder(nn.Module):
         x = self.pre_block(x)
         return self.block(x, rope)
 
-    def forward_features(self, x: torch.Tensor, rope: torch.Tensor) -> list[torch.Tensor]:
+    def forward_features(
+        self, x: torch.Tensor, rope: torch.Tensor, out_indices: Optional[list[int]] = None
+    ) -> list[torch.Tensor]:
         x = self.pre_block(x)
 
+        out_indices_set = set(out_indices) if out_indices is not None else None
         xs = []
-        for blk in self.block:
+        for idx, blk in enumerate(self.block):
             x = blk(x, rope)
-            xs.append(x)
+            if out_indices_set is None or idx in out_indices_set:
+                xs.append(x)
 
         return xs
 
@@ -438,6 +443,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         norm_layer_eps: float = self.config.get("norm_layer_eps", 1e-6)
         mlp_layer_type: str = self.config.get("mlp_layer_type", "FFN")
         act_layer_type: Optional[str] = self.config.get("act_layer_type", None)  # Default according to mlp type
+        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
         rope_rot_type: Literal["standard", "interleaved"] = self.config.get("rope_rot_type", "standard")
         rope_grid_indexing: Literal["ij", "xy"] = self.config.get("rope_grid_indexing", "ij")
         rope_grid_offset: int = self.config.get("rope_grid_offset", 0)
@@ -479,6 +485,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         self.norm_layer_eps = norm_layer_eps
         self.mlp_layer = mlp_layer
         self.act_layer = act_layer
+        self.out_indices = normalize_out_indices(out_indices, num_layers)
         self.rope_rot_type = rope_rot_type
         self.rope_grid_indexing = rope_grid_indexing
         self.rope_grid_offset = rope_grid_offset
@@ -571,8 +578,9 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
 
             self.attn_pool = MultiHeadAttentionPool(hidden_dim, attn_pool_num_heads, mlp_dim, qkv_bias=True)
 
-        self.return_stages = ["neck"]  # Actually meaningless, just for completeness
-        self.return_channels = [hidden_dim]
+        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.embedding_size = hidden_dim
         self.classifier = self.create_classifier()
 
@@ -658,8 +666,12 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)
 
+    def transform_to_backbone(self) -> None:
+        super().transform_to_backbone()
+        self.norm = nn.Identity()
+
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
         x = self.conv_proj(x)
         x = self.patch_embed(x)
 
@@ -679,15 +691,21 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         if self.pos_embed_special_tokens is True:
             x = x + self._get_pos_embed(H, W)
 
-        x = self.encoder(x, self._get_rope_embed(H, W))
-        x = self.norm(x)
+        rope = self._get_rope_embed(H, W)
+        if self.out_indices is None:
+            xs = [self.encoder(x, rope)]
+        else:
+            xs = self.encoder.forward_features(x, rope, out_indices=self.out_indices)
 
-        x = x[:, self.num_special_tokens :]
-        x = x.permute(0, 2, 1)
-        (B, C, _) = x.size()
-        x = x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+        out: dict[str, torch.Tensor] = {}
+        for stage_name, stage_x in zip(self.return_stages, xs):
+            stage_x = stage_x[:, self.num_special_tokens :]
+            stage_x = stage_x.permute(0, 2, 1)
+            B, C, _ = stage_x.size()
+            stage_x = stage_x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+            out[stage_name] = stage_x
 
-        return {self.return_stages[0]: x}
+        return out
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
@@ -709,7 +727,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         return_all_features: bool = False,
         return_keys: Literal["all", "tokens", "embedding"] = "tokens",
     ) -> TokenOmissionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)
@@ -789,7 +807,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         mask_token: Optional[torch.Tensor] = None,
         return_keys: Literal["all", "features", "embedding"] = "features",
     ) -> TokenRetentionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         x = self.conv_proj(x)
         x = mask_tensor(x, mask, mask_token=mask_token, patch_factor=self.max_stride // self.stem_stride)
@@ -820,7 +838,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         if return_keys in ("all", "features"):
             features = x[:, self.num_special_tokens :]
             features = features.permute(0, 2, 1)
-            (B, C, _) = features.size()
+            B, C, _ = features.size()
             features = features.reshape(B, C, H // self.patch_size, W // self.patch_size)
             result["features"] = features
 
@@ -840,7 +858,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         return result
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)

birder/net/sequencer2d.py

@@ -57,16 +57,16 @@ class LSTM2d(nn.Module):
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, H, W, C) = x.shape
+        B, H, W, C = x.shape
 
         v = x.permute(0, 2, 1, 3)
         v = v.reshape(-1, H, C)
-        (v, _) = self.rnn_v(v)
+        v, _ = self.rnn_v(v)
         v = v.reshape(B, W, H, -1)
         v = v.permute(0, 2, 1, 3)
 
         h = x.reshape(-1, W, C)
-        (h, _) = self.rnn_h(h)
+        h, _ = self.rnn_h(h)
         h = h.reshape(B, H, W, -1)
 
         x = torch.concat([v, h], dim=-1)
@@ -187,7 +187,6 @@ class Sequencer2d(BaseNet):
                 kernel_size=(patch_sizes[0], patch_sizes[0]),
                 stride=(patch_sizes[0], patch_sizes[0]),
                 padding=(0, 0),
-                bias=True,
             ),
             Permute([0, 2, 3, 1]),
         )

birder/net/shufflenet_v1.py

@@ -22,7 +22,7 @@ from birder.net.base import DetectorBackbone
 
 
 def channel_shuffle(x: torch.Tensor, groups: int) -> torch.Tensor:
-    (batch_size, num_channels, height, width) = x.size()
+    batch_size, num_channels, height, width = x.size()
     channels_per_group = num_channels // groups
 
     # Reshape

birder/net/shufflenet_v2.py

@@ -85,7 +85,7 @@ class ShuffleUnit(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.dw_conv_stride == 1:
-            (branch1, branch2) = x.chunk(2, dim=1)
+            branch1, branch2 = x.chunk(2, dim=1)
             x = torch.concat((branch1, self.branch2(branch2)), dim=1)
         else:
             x = torch.concat((self.branch1(x), self.branch2(x)), dim=1)

birder/net/simple_vit.py

@@ -26,17 +26,19 @@ 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 MaskedTokenOmissionMixin
 from birder.net.base import PreTrainEncoder
 from birder.net.base import TokenOmissionResultType
+from birder.net.base import normalize_out_indices
 from birder.net.base import pos_embedding_sin_cos_2d
 from birder.net.vit import Encoder
 from birder.net.vit import EncoderBlock
 from birder.net.vit import PatchEmbed
 
 
-# pylint: disable=invalid-name
-class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
+# pylint: disable=invalid-name,too-many-instance-attributes
+class Simple_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin):
     block_group_regex = r"encoder\.block\.(\d+)"
 
     def __init__(
@@ -56,6 +58,7 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
         num_heads: int = self.config["num_heads"]
         hidden_dim: int = self.config["hidden_dim"]
         mlp_dim: int = self.config["mlp_dim"]
+        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
         drop_path_rate: float = self.config["drop_path_rate"]
 
         torch._assert(image_size[0] % patch_size == 0, "Input shape indivisible by patch size!")
@@ -66,6 +69,7 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
         self.hidden_dim = hidden_dim
         self.mlp_dim = mlp_dim
         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.conv_proj = nn.Conv2d(
@@ -74,7 +78,6 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
             kernel_size=(patch_size, patch_size),
             stride=(patch_size, patch_size),
             padding=(0, 0),
-            bias=True,
         )
         self.patch_embed = PatchEmbed()
 
@@ -94,6 +97,9 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
             nn.Flatten(1),
         )
 
+        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.embedding_size = hidden_dim
         self.classifier = self.create_classifier()
 
@@ -144,7 +150,7 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
         return_all_features: bool = False,
         return_keys: Literal["all", "tokens", "embedding"] = "tokens",
     ) -> TokenOmissionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)
@@ -179,7 +185,7 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
         return result
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
         x = self.conv_proj(x)
         x = self.patch_embed(x)
         x = x + self._get_pos_embed(H, W)
@@ -193,6 +199,42 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
         x = x.permute(0, 2, 1)
         return self.features(x)
 
+    def transform_to_backbone(self) -> None:
+        super().transform_to_backbone()
+        self.norm = nn.Identity()
+
+    def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
+        H, W = x.shape[-2:]
+        x = self.conv_proj(x)
+        x = self.patch_embed(x)
+        x = x + self._get_pos_embed(H, W)
+
+        if self.out_indices is None:
+            xs = [self.encoder(x)]
+        else:
+            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):
+            stage_x = stage_x[:, self.num_special_tokens :]
+            stage_x = stage_x.permute(0, 2, 1)
+            B, C, _ = stage_x.size()
+            stage_x = stage_x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+            out[stage_name] = stage_x
+
+        return out
+
+    def freeze_stages(self, up_to_stage: int) -> None:
+        for param in self.conv_proj.parameters():
+            param.requires_grad_(False)
+
+        for idx, module in enumerate(self.encoder.children()):
+            if idx >= up_to_stage:
+                break
+
+            for param in module.parameters():
+                param.requires_grad_(False)
+
     def adjust_size(self, new_size: tuple[int, int]) -> None:
         if new_size == self.size:
             return