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

birder/introspection/transformer_attribution.py

@@ -66,7 +66,7 @@ def compute_attribution_rollout(
 
     mask = mask / (mask.max() + 1e-8)
 
-    (grid_h, grid_w) = patch_grid_shape
+    grid_h, grid_w = patch_grid_shape
     mask = mask.reshape(grid_h, grid_w)
 
     return mask
@@ -140,7 +140,7 @@ class TransformerAttribution:
         self.gatherer = AttributionGatherer(net, attention_layer_name)
 
     def __call__(self, image: str | Path | Image.Image, target_class: Optional[int] = None) -> InterpretabilityResult:
-        (input_tensor, rgb_img) = preprocess_image(image, self.transform, self.device)
+        input_tensor, rgb_img = preprocess_image(image, self.transform, self.device)
         input_tensor.requires_grad_(True)
 
         self.net.zero_grad()
@@ -156,7 +156,7 @@ class TransformerAttribution:
 
         attribution_data = self.gatherer.get_captured_data()
 
-        (_, _, H, W) = input_tensor.shape
+        _, _, H, W = input_tensor.shape
         patch_grid_shape = (H // self.net.stem_stride, W // self.net.stem_stride)
 
         attribution_map = compute_attribution_rollout(

birder/layers/attention_pool.py

@@ -39,13 +39,13 @@ class MultiHeadAttentionPool(nn.Module):
         nn.init.trunc_normal_(self.latent, std=dim**-0.5)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, N, C) = x.size()
+        B, N, C = x.size()
 
         q_latent = self.latent.expand(B, self.latent_len, -1)
         q = self.q(q_latent).reshape(B, self.latent_len, self.num_heads, self.head_dim).transpose(1, 2)
 
         kv = self.kv(x).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
-        (k, v) = kv.unbind(0)
+        k, v = kv.unbind(0)
 
         x = F.scaled_dot_product_attention(q, k, v, scale=self.scale)  # pylint: disable=not-callable
         x = x.transpose(1, 2).reshape(B, self.latent_len, C)

birder/model_registry/model_registry.py

@@ -6,6 +6,7 @@ from typing import Any
 from typing import Literal
 from typing import Optional
 
+from birder.conf.settings import DEFAULT_NUM_CHANNELS
 from birder.model_registry import manifest
 
 if TYPE_CHECKING is True:
@@ -229,8 +230,8 @@ class ModelRegistry:
     def net_factory(
         self,
         name: str,
-        input_channels: int,
         num_classes: int,
+        input_channels: int = DEFAULT_NUM_CHANNELS,
         *,
         config: Optional[dict[str, Any]] = None,
         size: Optional[tuple[int, int]] = None,

birder/net/__init__.py

@@ -6,6 +6,7 @@ from birder.net.coat import CoaT
 from birder.net.conv2former import Conv2Former
 from birder.net.convmixer import ConvMixer
 from birder.net.convnext_v1 import ConvNeXt_v1
+from birder.net.convnext_v1_iso import ConvNeXt_v1_Isotropic
 from birder.net.convnext_v2 import ConvNeXt_v2
 from birder.net.crossformer import CrossFormer
 from birder.net.crossvit import CrossViT
@@ -118,6 +119,7 @@ __all__ = [
     "Conv2Former",
     "ConvMixer",
     "ConvNeXt_v1",
+    "ConvNeXt_v1_Isotropic",
     "ConvNeXt_v2",
     "CrossFormer",
     "CrossViT",

birder/net/_rope_vit_configs.py

@@ -88,6 +88,11 @@ def register_rope_vit_configs(rope_vit: type[BaseNet]) -> None:
         rope_vit,
         config={"patch_size": 16, **SMALL},
     )
+    registry.register_model_config(
+        "rope_vit_s16_avg",
+        rope_vit,
+        config={"patch_size": 16, **SMALL, "class_token": False},
+    )
     registry.register_model_config(
         "rope_i_vit_s16_pn_aps_c1",  # For PE Core - https://arxiv.org/abs/2504.13181
         rope_vit,

birder/net/_vit_configs.py

@@ -215,19 +215,6 @@ def register_vit_configs(vit: type[BaseNet]) -> None:
             "drop_path_rate": 0.1,
         },
     )
-    registry.register_model_config(  # From "Scaling Vision Transformers to 22 Billion Parameters"
-        "vit_22b_p16_qkn",
-        vit,
-        config={
-            "patch_size": 16,
-            "num_layers": 48,
-            "num_heads": 48,
-            "hidden_dim": 6144,
-            "mlp_dim": 24576,
-            "qk_norm": True,
-            "drop_path_rate": 0.1,
-        },
-    )
 
     # With registers
     ####################

birder/net/alexnet.py

@@ -27,17 +27,17 @@ class AlexNet(BaseNet):
         assert self.config is None, "config not supported"
 
         self.body = nn.Sequential(
-            nn.Conv2d(self.input_channels, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2), bias=True),
+            nn.Conv2d(self.input_channels, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2)),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)),
-            nn.Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), bias=True),
+            nn.Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)),
-            nn.Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=True),
+            nn.Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
             nn.ReLU(inplace=True),
-            nn.Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=True),
+            nn.Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
             nn.ReLU(inplace=True),
-            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=True),
+            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)),
             nn.AdaptiveAvgPool2d(output_size=(6, 6)),

birder/net/base.py

@@ -5,6 +5,7 @@ from typing import Literal
 from typing import NotRequired
 from typing import Optional
 from typing import TypedDict
+from typing import overload
 
 import torch
 import torch.nn.functional as F
@@ -54,6 +55,30 @@ def make_divisible(v: float, divisor: int, min_value: Optional[int] = None) -> i
     return new_v
 
 
+@overload
+def normalize_out_indices(out_indices: None, num_layers: int) -> None: ...
+
+
+@overload
+def normalize_out_indices(out_indices: list[int], num_layers: int) -> list[int]: ...
+
+
+def normalize_out_indices(out_indices: Optional[list[int]], num_layers: int) -> Optional[list[int]]:
+    if out_indices is None:
+        return None
+
+    normalized_indices = []
+    for idx in out_indices:
+        if idx < 0:
+            idx = num_layers + idx
+        if idx < 0 or idx >= num_layers:
+            raise ValueError(f"out_indices contains invalid index for num_layers={num_layers}")
+
+        normalized_indices.append(idx)
+
+    return normalized_indices
+
+
 # class MiscNet(nn.Module):
 #     """
 #     Base class for general-purpose neural networks with automatic model registration
@@ -137,8 +162,8 @@ class BaseNet(nn.Module):
 
         self.dynamic_size = False
 
-        self.classifier: nn.Module
         self.embedding_size: int
+        self.classifier: nn.Module
 
     def create_classifier(self, embed_dim: Optional[int] = None) -> nn.Module:
         if self.num_classes == 0:
@@ -274,7 +299,7 @@ def pos_embedding_sin_cos_2d(
 ) -> torch.Tensor:
     # assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sin-cos emb"
 
-    (y, x) = torch.meshgrid(torch.arange(h, device=device), torch.arange(w, device=device), indexing="ij")
+    y, x = torch.meshgrid(torch.arange(h, device=device), torch.arange(w, device=device), indexing="ij")
     omega = torch.arange(dim // 4, device=device) / (dim // 4 - 1)
     omega = 1.0 / (temperature**omega)
 
@@ -294,7 +319,7 @@ def interpolate_attention_bias(
     new_resolution: tuple[int, int],
     mode: Literal["bilinear", "bicubic"] = "bicubic",
 ) -> torch.Tensor:
-    (H, _) = attention_bias.size()
+    H, _ = attention_bias.size()
 
     # Interpolate
     orig_dtype = attention_bias.dtype

birder/net/biformer.py

@@ -30,7 +30,7 @@ from birder.net.base import DetectorBackbone
 
 
 def _grid2seq(x: torch.Tensor, region_size: tuple[int, int], num_heads: int) -> tuple[torch.Tensor, int, int]:
-    (B, C, H, W) = x.size()
+    B, C, H, W = x.size()
     region_h = H // region_size[0]
     region_w = W // region_size[1]
     x = x.view(B, num_heads, C // num_heads, region_h, region_size[0], region_w, region_size[1])
@@ -40,7 +40,7 @@ def _grid2seq(x: torch.Tensor, region_size: tuple[int, int], num_heads: int) ->
 
 
 def _seq2grid(x: torch.Tensor, region_h: int, region_w: int, region_size: tuple[int, int]) -> torch.Tensor:
-    (bs, n_head, _, _, head_dim) = x.size()
+    bs, n_head, _, _, head_dim = x.size()
     x = x.view(bs, n_head, region_h, region_w, region_size[0], region_size[1], head_dim)
     x = torch.einsum("bmhwpqd->bmdhpwq", x).reshape(
         bs, n_head * head_dim, region_h * region_size[0], region_w * region_size[1]
@@ -60,7 +60,7 @@ def regional_routing_attention_torch(
     auto_pad: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     kv_region_size = region_size
-    (bs, n_head, q_nregion, topk) = region_graph.size()
+    bs, n_head, q_nregion, topk = region_graph.size()
 
     # Pad to deal with any input size
     q_pad_b = 0
@@ -68,13 +68,13 @@ def regional_routing_attention_torch(
     kv_pad_b = 0
     kv_pad_r = 0
     if auto_pad is True:
-        (_, _, h_q, w_q) = query.size()
+        _, _, h_q, w_q = query.size()
         q_pad_b = (region_size[0] - h_q % region_size[0]) % region_size[0]
         q_pad_r = (region_size[1] - w_q % region_size[1]) % region_size[1]
         if q_pad_b > 0 or q_pad_r > 0:
             query = F.pad(query, (0, q_pad_r, 0, q_pad_b))
 
-        (_, _, h_k, w_k) = key.size()
+        _, _, h_k, w_k = key.size()
         kv_pad_b = (kv_region_size[0] - h_k % kv_region_size[0]) % kv_region_size[0]
         kv_pad_r = (kv_region_size[1] - w_k % kv_region_size[1]) % kv_region_size[1]
         if kv_pad_r > 0 or kv_pad_b > 0:
@@ -87,12 +87,12 @@ def regional_routing_attention_torch(
         w_k = None
 
     # To sequence format
-    (query, q_region_h, q_region_w) = _grid2seq(query, region_size=region_size, num_heads=n_head)
-    (key, _, _) = _grid2seq(key, region_size=kv_region_size, num_heads=n_head)
-    (value, _, _) = _grid2seq(value, region_size=kv_region_size, num_heads=n_head)
+    query, q_region_h, q_region_w = _grid2seq(query, region_size=region_size, num_heads=n_head)
+    key, _, _ = _grid2seq(key, region_size=kv_region_size, num_heads=n_head)
+    value, _, _ = _grid2seq(value, region_size=kv_region_size, num_heads=n_head)
 
     # Gather key and values
-    (bs, n_head, kv_nregion, kv_region_size, head_dim) = key.size()
+    bs, n_head, kv_nregion, kv_region_size, head_dim = key.size()
     broadcasted_region_graph = region_graph.view(bs, n_head, q_nregion, topk, 1, 1).expand(
         -1, -1, -1, -1, kv_region_size, head_dim
     )
@@ -146,12 +146,12 @@ class BiLevelRoutingAttention(nn.Module):
         self.output_linear = nn.Conv2d(dim, dim, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (_, _, H, W) = x.size()
+        _, _, H, W = x.size()
         region_size = (H // self.n_win_h, W // self.n_win_w)
 
         # Linear projection
         qkv = self.qkv_linear(x)
-        (q, k, v) = qkv.chunk(3, dim=1)
+        q, k, v = qkv.chunk(3, dim=1)
 
         # Region-to-region routing
         q_r = F.avg_pool2d(  # pylint: disable=not-callable
@@ -163,11 +163,11 @@ class BiLevelRoutingAttention(nn.Module):
         q_r = q_r.permute(0, 2, 3, 1).flatten(1, 2)  # (n, (hw), c)
         k_r = k_r.flatten(2, 3)  # (n, c, (hw))
         a_r = q_r @ k_r
-        (_, idx_r) = torch.topk(a_r, k=self.topk, dim=-1)
+        _, idx_r = torch.topk(a_r, k=self.topk, dim=-1)
         idx_r = idx_r.unsqueeze_(1).expand(-1, self.num_heads, -1, -1)
 
         # Token to token attention
-        (output, _) = regional_routing_attention_torch(
+        output, _ = regional_routing_attention_torch(
             q, k, v, scale=self.scale, region_graph=idx_r, region_size=region_size, auto_pad=True
         )
 
@@ -190,12 +190,12 @@ class Attention(nn.Module):
         self.proj_drop = nn.Dropout(proj_drop)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, C, H, W) = x.size()
+        B, C, H, W = x.size()
         x = x.permute(0, 2, 3, 1).reshape(B, H * W, C)
 
         N = H * W
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        (q, k, v) = qkv.unbind(0)
+        q, k, v = qkv.unbind(0)
 
         x = F.scaled_dot_product_attention(  # pylint: disable=not-callable
             q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0, scale=self.scale
@@ -237,8 +237,8 @@ class AttentionLePE(nn.Module):
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, C, H, W) = x.size()
-        (q, k, v) = self.qkv(x).chunk(3, dim=1)
+        B, C, H, W = x.size()
+        q, k, v = self.qkv(x).chunk(3, dim=1)
 
         attn = q.view(B, self.num_heads, self.head_dim, H * W).transpose(-1, -2) @ k.view(
             B, self.num_heads, self.head_dim, H * W

birder/net/cait.py

@@ -47,7 +47,7 @@ class ClassAttention(nn.Module):
         self.proj_drop = nn.Dropout(proj_drop)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, N, C) = x.shape
+        B, N, C = x.shape
         q = self.q(x[:, 0]).unsqueeze(1).reshape(B, 1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
         k = self.k(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
         v = self.v(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
@@ -103,7 +103,7 @@ class TalkingHeadAttn(nn.Module):
         self.proj_drop = nn.Dropout(proj_drop)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, N, C) = x.shape
+        B, N, C = x.shape
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
         q = qkv[0] * self.scale
         k = qkv[1]

birder/net/cas_vit.py

@@ -122,7 +122,7 @@ class AdditiveTokenMixer(nn.Module):
         self.proj_drop = nn.Dropout(proj_drop)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (q, k, v) = self.qkv(x).chunk(3, dim=1)
+        q, k, v = self.qkv(x).chunk(3, dim=1)
         q = self.op_q(q)
         k = self.op_k(k)
 

birder/net/coat.py

@@ -57,8 +57,8 @@ class ConvRelPosEnc(nn.Module):
         self.channel_splits = [x * head_channels for x in head_splits]
 
     def forward(self, q: torch.Tensor, v: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
-        (B, num_heads, N, C) = q.size()
-        (H, W) = size
+        B, num_heads, N, C = q.size()
+        H, W = size
         torch._assert(N == 1 + H * W, "size mismatch")  # pylint: disable=protected-access
 
         # Convolutional relative position encoding.
@@ -102,11 +102,11 @@ class FactorAttnConvRelPosEnc(nn.Module):
         self.crpe = shared_crpe
 
     def forward(self, x: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
-        (B, N, C) = x.size()
+        B, N, C = x.size()
 
         # Generate Q, K, V
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        (q, k, v) = qkv.unbind(0)  # [B, h, N, Ch]
+        q, k, v = qkv.unbind(0)  # [B, h, N, Ch]
 
         # Factorized attention
         k_softmax = k.softmax(dim=2)
@@ -135,8 +135,8 @@ class ConvPosEnc(nn.Module):
         )
 
     def forward(self, x: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
-        (B, N, C) = x.size()
-        (H, W) = size
+        B, N, C = x.size()
+        H, W = size
         torch._assert(N == 1 + H * W, "size mismatch")  # pylint: disable=protected-access
 
         # Extract CLS token and image tokens
@@ -244,8 +244,8 @@ class ParallelBlock(nn.Module):
         return self.interpolate(x, scale_factor=1.0 / factor, size=size)
 
     def interpolate(self, x: torch.Tensor, scale_factor: float, size: tuple[int, int]) -> torch.Tensor:
-        (B, N, C) = x.size()
-        (H, W) = size
+        B, N, C = x.size()
+        H, W = size
         torch._assert(N == 1 + H * W, "size mismatch")  # pylint: disable=protected-access
 
         cls_token = x[:, :1, :]
@@ -268,7 +268,7 @@ class ParallelBlock(nn.Module):
     def forward(
         self, x1: torch.Tensor, x2: torch.Tensor, x3: torch.Tensor, x4: torch.Tensor, sizes: list[tuple[int, int]]
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        (_, s2, s3, s4) = sizes
+        _, s2, s3, s4 = sizes
         cur2 = self.norm12(x2)
         cur3 = self.norm13(x3)
         cur4 = self.norm14(x4)
@@ -310,7 +310,7 @@ class PatchEmbed(nn.Module):
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int]]:
         x = self.proj(x)
-        (H, W) = x.shape[2:4]
+        H, W = x.shape[2:4]
 
         x = x.flatten(2).transpose(1, 2)
         x = self.norm(x)
@@ -500,7 +500,7 @@ class CoaT(DetectorBackbone):
         B = x.shape[0]
 
         # Serial blocks 1
-        (x1, (h1, w1)) = self.patch_embed1(x)
+        x1, (h1, w1) = self.patch_embed1(x)
         x1 = insert_cls(x1, self.cls_token1)
         for blk in self.serial_blocks1:
             x1 = blk(x1, size=(h1, w1))
@@ -508,7 +508,7 @@ class CoaT(DetectorBackbone):
         x1_no_cls = remove_cls(x1).reshape(B, h1, w1, -1).permute(0, 3, 1, 2).contiguous()
 
         # Serial blocks 2
-        (x2, (h2, w2)) = self.patch_embed2(x1_no_cls)
+        x2, (h2, w2) = self.patch_embed2(x1_no_cls)
         x2 = insert_cls(x2, self.cls_token2)
         for blk in self.serial_blocks2:
             x2 = blk(x2, size=(h2, w2))
@@ -516,7 +516,7 @@ class CoaT(DetectorBackbone):
         x2_no_cls = remove_cls(x2).reshape(B, h2, w2, -1).permute(0, 3, 1, 2).contiguous()
 
         # Serial blocks 3
-        (x3, (h3, w3)) = self.patch_embed3(x2_no_cls)
+        x3, (h3, w3) = self.patch_embed3(x2_no_cls)
         x3 = insert_cls(x3, self.cls_token3)
         for blk in self.serial_blocks3:
             x3 = blk(x3, size=(h3, w3))
@@ -524,7 +524,7 @@ class CoaT(DetectorBackbone):
         x3_no_cls = remove_cls(x3).reshape(B, h3, w3, -1).permute(0, 3, 1, 2).contiguous()
 
         # Serial blocks 4
-        (x4, (h4, w4)) = self.patch_embed4(x3_no_cls)
+        x4, (h4, w4) = self.patch_embed4(x3_no_cls)
         x4 = insert_cls(x4, self.cls_token4)
         for blk in self.serial_blocks4:
             x4 = blk(x4, size=(h4, w4))
@@ -537,7 +537,7 @@ class CoaT(DetectorBackbone):
                 x2 = self.cpe2(x2, (h2, w2))
                 x3 = self.cpe3(x3, (h3, w3))
                 x4 = self.cpe4(x4, (h4, w4))
-                (x1, x2, x3, x4) = blk(x1, x2, x3, x4, sizes=[(h1, w1), (h2, w2), (h3, w3), (h4, w4)])
+                x1, x2, x3, x4 = blk(x1, x2, x3, x4, sizes=[(h1, w1), (h2, w2), (h3, w3), (h4, w4)])
 
             x1_no_cls = remove_cls(x1).reshape(B, h1, w1, -1).permute(0, 3, 1, 2).contiguous()
             x2_no_cls = remove_cls(x2).reshape(B, h2, w2, -1).permute(0, 3, 1, 2).contiguous()

birder/net/convnext_v1.py

@@ -37,15 +37,7 @@ class ConvNeXtBlock(nn.Module):
     ) -> None:
         super().__init__()
         self.block = nn.Sequential(
-            nn.Conv2d(
-                channels,
-                channels,
-                kernel_size=(7, 7),
-                stride=(1, 1),
-                padding=(3, 3),
-                groups=channels,
-                bias=True,
-            ),
+            nn.Conv2d(channels, channels, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=channels),
             Permute([0, 2, 3, 1]),
             nn.LayerNorm(channels, eps=1e-6),
             nn.Linear(channels, 4 * channels),  # Same as 1x1 conv
@@ -119,7 +111,7 @@ class ConvNeXt_v1(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
                 layers.append(
                     nn.Sequential(
                         LayerNorm2d(i, eps=1e-6),
-                        nn.Conv2d(i, out, kernel_size=(2, 2), stride=(2, 2), padding=(0, 0), bias=True),
+                        nn.Conv2d(i, out, kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)),
                     )
                 )
 

birder/net/convnext_v1_iso.py

@@ -0,0 +1,198 @@
+"""
+ConvNeXt v1 Isotropic, adapted from
+https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext_isotropic.py
+
+Paper "A ConvNet for the 2020s", https://arxiv.org/abs/2201.03545
+"""
+
+# Reference license: MIT
+
+from functools import partial
+from typing import Any
+from typing import Literal
+from typing import Optional
+
+import torch
+from torch import nn
+from torchvision.ops import Permute
+from torchvision.ops import StochasticDepth
+
+from birder.common.masking import mask_tensor
+from birder.layers import LayerNorm2d
+from birder.model_registry import registry
+from birder.net.base import DetectorBackbone
+from birder.net.base import MaskedTokenRetentionMixin
+from birder.net.base import PreTrainEncoder
+from birder.net.base import TokenRetentionResultType
+from birder.net.base import normalize_out_indices
+
+
+class ConvNeXtBlock(nn.Module):
+    def __init__(self, channels: int, stochastic_depth_prob: float) -> None:
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(channels, channels, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=channels),
+            Permute([0, 2, 3, 1]),
+            nn.LayerNorm(channels, eps=1e-6),
+            nn.Linear(channels, 4 * channels),
+            nn.GELU(),
+            nn.Linear(4 * channels, channels),
+            Permute([0, 3, 1, 2]),
+        )
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, mode="row")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        x = self.block(x)
+        x = self.stochastic_depth(x)
+        x += identity
+
+        return x
+
+
+# pylint: disable=invalid-name
+class ConvNeXt_v1_Isotropic(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
+    block_group_regex = r"body\.(\d+)"
+
+    def __init__(
+        self,
+        input_channels: int,
+        num_classes: int,
+        *,
+        config: Optional[dict[str, Any]] = None,
+        size: Optional[tuple[int, int]] = None,
+    ) -> None:
+        super().__init__(input_channels, num_classes, config=config, size=size)
+        assert self.config is not None, "must set config"
+
+        patch_size = 16
+        dim: int = self.config["dim"]
+        num_layers: int = self.config["num_layers"]
+        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
+        drop_path_rate: float = self.config["drop_path_rate"]
+
+        torch._assert(self.size[0] % patch_size == 0, "Input shape indivisible by patch size!")
+        torch._assert(self.size[1] % patch_size == 0, "Input shape indivisible by patch size!")
+        self.patch_size = patch_size
+        self.out_indices = normalize_out_indices(out_indices, num_layers)
+
+        self.stem = nn.Conv2d(
+            self.input_channels,
+            dim,
+            kernel_size=(patch_size, patch_size),
+            stride=(patch_size, patch_size),
+            padding=(0, 0),
+        )
+
+        layers = []
+        for idx in range(num_layers):
+            # Adjust stochastic depth probability based on the depth of the stage block
+            sd_prob = drop_path_rate * idx / (num_layers - 1.0)
+            layers.append(ConvNeXtBlock(dim, sd_prob))
+
+        self.body = nn.Sequential(*layers)
+        self.features = nn.Sequential(
+            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
+            LayerNorm2d(dim, eps=1e-6),
+            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 = [dim] * num_return_stages
+        self.embedding_size = dim
+        self.classifier = self.create_classifier()
+
+        self.max_stride = patch_size
+        self.stem_stride = patch_size
+        self.stem_width = dim
+        self.encoding_size = dim
+        self.decoder_block = partial(ConvNeXtBlock, stochastic_depth_prob=0)
+
+        # Weights initialization
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
+        x = self.stem(x)
+
+        if self.out_indices is None:
+            x = self.body(x)
+            return {self.return_stages[0]: x}
+
+        stage_num = 0
+        out: dict[str, torch.Tensor] = {}
+        for idx, module in enumerate(self.body.children()):
+            x = module(x)
+            if idx in self.out_indices:
+                out[self.return_stages[stage_num]] = x
+                stage_num += 1
+
+        return out
+
+    def freeze_stages(self, up_to_stage: int) -> None:
+        for param in self.stem.parameters():
+            param.requires_grad_(False)
+
+        for idx, module in enumerate(self.body.children()):
+            if idx >= up_to_stage:
+                break
+
+            for param in module.parameters():
+                param.requires_grad_(False)
+
+    def masked_encoding_retention(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        mask_token: Optional[torch.Tensor] = None,
+        return_keys: Literal["all", "features", "embedding"] = "features",
+    ) -> TokenRetentionResultType:
+        x = self.stem(x)
+        x = mask_tensor(x, mask, patch_factor=self.max_stride // self.stem_stride, mask_token=mask_token)
+        x = self.body(x)
+
+        result: TokenRetentionResultType = {}
+        if return_keys in ("all", "features"):
+            result["features"] = x
+        if return_keys in ("all", "embedding"):
+            result["embedding"] = self.features(x)
+
+        return result
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.stem(x)
+        return self.body(x)
+
+    def embedding(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.forward_features(x)
+        return self.features(x)
+
+    def adjust_size(self, new_size: tuple[int, int]) -> None:
+        if new_size == self.size:
+            return
+
+        assert new_size[0] % self.patch_size == 0, "Input shape indivisible by patch size!"
+        assert new_size[1] % self.patch_size == 0, "Input shape indivisible by patch size!"
+
+        super().adjust_size(new_size)
+
+
+registry.register_model_config(
+    "convnext_v1_iso_small",
+    ConvNeXt_v1_Isotropic,
+    config={"dim": 384, "num_layers": 18, "drop_path_rate": 0.1},
+)
+registry.register_model_config(
+    "convnext_v1_iso_base",
+    ConvNeXt_v1_Isotropic,
+    config={"in_channels": 768, "num_layers": 18, "drop_path_rate": 0.2},
+)
+registry.register_model_config(
+    "convnext_v1_iso_large",
+    ConvNeXt_v1_Isotropic,
+    config={"in_channels": 1024, "num_layers": 36, "drop_path_rate": 0.5},
+)