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

birder/net/cswin_transformer.py

@@ -29,7 +29,7 @@ from birder.net.vit import PatchEmbed
 
 
 def img2windows(img: torch.Tensor, h_sp: int, w_sp: int) -> torch.Tensor:
-    (B, C, H, W) = img.size()
+    B, C, H, W = img.size()
     img_reshape = img.view(B, C, H // h_sp, h_sp, W // w_sp, w_sp)
     img_perm = img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, h_sp * w_sp, C)
 
@@ -81,7 +81,7 @@ class LePEAttention(nn.Module):
             raise ValueError("unsupported idx")
 
     def im2cswin(self, x: torch.Tensor) -> torch.Tensor:
-        (B, _, C) = x.size()
+        B, _, C = x.size()
         x = x.transpose(-2, -1).contiguous().view(B, C, self.resolution[0], self.resolution[1])
         x = img2windows(x, self.h_sp, self.w_sp)
         x = x.reshape(-1, self.h_sp * self.w_sp, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous()
@@ -89,7 +89,7 @@ class LePEAttention(nn.Module):
         return x
 
     def get_lepe(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        (B, _, C) = x.size()
+        B, _, C = x.size()
         H = self.resolution[0]
         W = self.resolution[1]
         x = x.transpose(-2, -1).contiguous().view(B, C, H, W)
@@ -107,13 +107,13 @@ class LePEAttention(nn.Module):
         return (x, lepe)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (q, k, v) = x.unbind(0)
+        q, k, v = x.unbind(0)
 
-        (B, _, C) = q.shape
+        B, _, C = q.shape
 
         q = self.im2cswin(q)
         k = self.im2cswin(k)
-        (v, lepe) = self.get_lepe(v)
+        v, lepe = self.get_lepe(v)
 
         q = q * self.scale
         attn = q @ k.transpose(-2, -1)  # B head N C @ B head C N --> B head N N
@@ -136,12 +136,12 @@ class MergeBlock(nn.Module):
         self.resolution = resolution
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, _, C) = x.size()
+        B, _, C = x.size()
         H = self.resolution[0]
         W = self.resolution[1]
         x = x.transpose(-2, -1).contiguous().view(B, C, H, W)
         x = self.conv(x)
-        (B, C) = x.shape[:2]
+        B, C = x.shape[:2]
         x = x.view(B, C, -1).transpose(-2, -1).contiguous()
         x = self.norm(x)
 
@@ -206,7 +206,7 @@ class CSWinBlock(nn.Module):
         self.drop_path = StochasticDepth(drop_path, mode="row")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, _, C) = x.shape
+        B, _, C = x.shape
 
         qkv = self.qkv(self.norm1(x)).reshape(B, -1, 3, C).permute(2, 0, 1, 3)
         if self.branch_num == 2:
@@ -350,7 +350,7 @@ class CSWin_Transformer(DetectorBackbone):
         for name, module in self.body.named_children():
             x = module(x)
             if name in self.return_stages:
-                (B, L, C) = x.size()
+                B, L, C = x.size()
                 H = int(math.sqrt(L))
                 W = H
                 out[name] = x.transpose(-2, -1).contiguous().view(B, C, H, W)

birder/net/davit.py

@@ -31,7 +31,7 @@ from birder.net.base import TokenRetentionResultType
 
 
 def window_partition(x: torch.Tensor, window_size: tuple[int, int]) -> torch.Tensor:
-    (B, H, W, C) = x.shape
+    B, H, W, C = x.shape
     x = x.view(B, H // window_size[0], window_size[0], W // window_size[1], window_size[1], C)
     windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size[0], window_size[1], C)
 
@@ -64,7 +64,7 @@ class ConvPosEnc(nn.Module):
             dim,
             kernel_size=kernel_size,
             stride=(1, 1),
-            padding=(kernel_size[0] // 2, kernel_size[1] // 2),
+            padding=((kernel_size[0] - 1) // 2, (kernel_size[1] - 1) // 2),
             groups=dim,
         )
         if act is True:
@@ -92,10 +92,10 @@ class Downsample(nn.Module):
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (_, _, H, W) = x.shape
+        _, _, H, W = x.shape
         x = self.norm(x)
         if self.even_k is True:
-            (k_h, k_w) = self.conv.kernel_size
+            k_h, k_w = self.conv.kernel_size
             pad_r = (k_w - W % k_w) % k_w
             pad_b = (k_h - H % k_h) % k_h
             x = F.pad(x, (0, pad_r, 0, pad_b))
@@ -115,10 +115,10 @@ class ChannelAttention(nn.Module):
         self.proj = nn.Linear(dim, dim)
 
     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, k, v) = qkv.unbind(0)
+        q, k, v = qkv.unbind(0)
 
         k = k * self.scale
         attn = k.transpose(-1, -2) @ v
@@ -151,7 +151,7 @@ class ChannelBlock(nn.Module):
         self.drop_path = StochasticDepth(drop_path, mode="row")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, C, H, W) = x.shape
+        B, C, H, W = x.shape
         x = self.cpe1(x).flatten(2).transpose(1, 2)
 
         cur = self.norm1(x)
@@ -177,10 +177,10 @@ class WindowAttention(nn.Module):
         self.proj = nn.Linear(dim, dim)
 
     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, k, v) = qkv.unbind(0)
+        q, k, v = qkv.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, N, C)
@@ -215,7 +215,7 @@ class SpatialBlock(nn.Module):
 
     # pylint: disable=invalid-name
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, C, H, W) = x.shape
+        B, C, H, W = x.shape
 
         shortcut = self.cpe1(x).flatten(2).transpose(1, 2)
 
@@ -226,7 +226,7 @@ class SpatialBlock(nn.Module):
         pad_r = (self.window_size[1] - W % self.window_size[1]) % self.window_size[1]
         pad_b = (self.window_size[0] - H % self.window_size[0]) % self.window_size[0]
         x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
-        (_, Hp, Wp, _) = x.shape
+        _, Hp, Wp, _ = x.shape
 
         x_windows = window_partition(x, self.window_size)
         x_windows = x_windows.view(-1, self.window_size[0] * self.window_size[1], C)

birder/net/deit.py

@@ -16,13 +16,18 @@ import torch
 from torch import nn
 
 from birder.model_registry import registry
-from birder.net.base import BaseNet
+from birder.net._vit_configs import BASE
+from birder.net._vit_configs import SMALL
+from birder.net._vit_configs import TINY
+from birder.net.base import DetectorBackbone
+from birder.net.base import normalize_out_indices
 from birder.net.vit import Encoder
 from birder.net.vit import PatchEmbed
 from birder.net.vit import adjust_position_embedding
 
 
-class DeiT(BaseNet):
+# pylint: disable=too-many-instance-attributes
+class DeiT(DetectorBackbone):
     block_group_regex = r"encoder\.block\.(\d+)"
 
     def __init__(
@@ -44,6 +49,7 @@ class DeiT(BaseNet):
         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!")
@@ -53,6 +59,7 @@ class DeiT(BaseNet):
         self.num_layers = num_layers
         self.hidden_dim = hidden_dim
         self.num_special_tokens = 2
+        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(
@@ -61,7 +68,6 @@ class DeiT(BaseNet):
             kernel_size=(patch_size, patch_size),
             stride=(patch_size, patch_size),
             padding=(0, 0),
-            bias=True,
         )
         self.patch_embed = PatchEmbed()
 
@@ -89,11 +95,18 @@ class DeiT(BaseNet):
         )
         self.norm = nn.LayerNorm(hidden_dim, eps=1e-6)
 
+        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.dist_classifier = self.create_classifier()
         self.classifier = self.create_classifier()
         self.distillation_output = False
 
+        self.max_stride = patch_size
+        self.stem_stride = patch_size
+        self.stem_width = hidden_dim
+
         # Weight initialization
         if isinstance(self.conv_proj, nn.Conv2d):
             # Init the patchify stem
@@ -129,6 +142,53 @@ class DeiT(BaseNet):
     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()
+        self.dist_classifier = nn.Identity()
+
+    def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
+        H, W = x.shape[-2:]
+
+        # Reshape and permute the input tensor
+        x = self.conv_proj(x)
+        x = self.patch_embed(x)
+
+        # Expand the class token to the full batch
+        batch_class_token = self.class_token.expand(x.shape[0], -1, -1)
+        batch_dist_token = self.dist_token.expand(x.shape[0], -1, -1)
+
+        x = torch.concat([batch_class_token, batch_dist_token, x], dim=1)
+        x = x + self.pos_embedding
+
+        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)
+
+        self.pos_embedding.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 forward_features(self, x: torch.Tensor) -> torch.Tensor:
         # Reshape and permute the input tensor
         x = self.conv_proj(x)
@@ -200,38 +260,17 @@ class DeiT(BaseNet):
 registry.register_model_config(
     "deit_t16",
     DeiT,
-    config={
-        "patch_size": 16,
-        "num_layers": 12,
-        "num_heads": 3,
-        "hidden_dim": 192,
-        "mlp_dim": 768,
-        "drop_path_rate": 0.0,
-    },
+    config={"patch_size": 16, **TINY},
 )
 registry.register_model_config(
     "deit_s16",
     DeiT,
-    config={
-        "patch_size": 16,
-        "num_layers": 12,
-        "num_heads": 6,
-        "hidden_dim": 384,
-        "mlp_dim": 1536,
-        "drop_path_rate": 0.1,
-    },
+    config={"patch_size": 16, **SMALL, "drop_path_rate": 0.1},  # Override the SMALL definition
 )
 registry.register_model_config(
     "deit_b16",
     DeiT,
-    config={
-        "patch_size": 16,
-        "num_layers": 12,
-        "num_heads": 12,
-        "hidden_dim": 768,
-        "mlp_dim": 3072,
-        "drop_path_rate": 0.1,
-    },
+    config={"patch_size": 16, **BASE},
 )
 
 registry.register_weights(
@@ -242,7 +281,7 @@ registry.register_weights(
         "formats": {
             "pt": {
                 "file_size": 21.7,
-                "sha256": "ac124122dec9f1bceff383a6a555ca375ca1b613caf486dac3f29d87afac03b3",
+                "sha256": "68b33aba0c1be5e78d4a33e74a7c1ea72b6abb232d59f0048ff9b8342e43246e",
             }
         },
         "net": {"network": "deit_t16", "tag": "il-common"},
@@ -258,7 +297,7 @@ registry.register_weights(
         "formats": {
             "pt": {
                 "file_size": 21.7,
-                "sha256": "fafd0c3c65f9c35318f449f60485f640917736ee7b44056be55c2226909ffdb8",
+                "sha256": "f693e89fc350341141c55152bec9f499df63738e8423071f3b8e71801c3e5415",
             }
         },
         "net": {"network": "deit_t16", "tag": "dist-il-common"},