birder 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

birder/net/lit_v1_tiny.py

@@ -0,0 +1,342 @@
+"""
+LIT v1 Tiny, adapted from
+https://github.com/ziplab/LIT/blob/main/classification/code_for_lit_ti/lit.py
+
+Paper "Less is More: Pay Less Attention in Vision Transformers", https://arxiv.org/abs/2105.14217
+
+Generated by Claude Code Opus 4.5
+"""
+
+# Reference license: Apache-2.0
+
+import math
+from collections import OrderedDict
+from typing import Any
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torchvision.ops import Permute
+from torchvision.ops import StochasticDepth
+
+from birder.model_registry import registry
+from birder.net.base import DetectorBackbone
+from birder.net.lit_v1 import MLP
+from birder.net.lit_v1 import DeformablePatchMerging
+from birder.net.lit_v1 import IdentityDownsample
+from birder.net.vit import adjust_position_embedding
+
+
+class MLPBlock(nn.Module):
+    def __init__(self, dim: int, mlp_ratio: float, drop_path: float) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.mlp = MLP(dim, int(dim * mlp_ratio))
+        self.drop_path = StochasticDepth(drop_path, mode="row")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.drop_path(self.mlp(self.norm(x)))
+
+
+class Attention(nn.Module):
+    def __init__(self, dim: int, num_heads: int) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.scale = (dim // num_heads) ** -0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        (B, N, C) = x.size()
+        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)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = F.softmax(attn, dim=-1)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+
+        return x
+
+
+class ViTBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, mlp_ratio: float, drop_path: float) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
+        self.attn = Attention(dim, num_heads)
+        self.drop_path1 = StochasticDepth(drop_path, mode="row")
+
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
+        self.mlp = MLP(dim, int(dim * mlp_ratio))
+        self.drop_path2 = StochasticDepth(drop_path, mode="row")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.drop_path1(self.attn(self.norm1(x)))
+        x = x + self.drop_path2(self.mlp(self.norm2(x)))
+
+        return x
+
+
+class LITStage(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        input_resolution: tuple[int, int],
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        has_msa: bool,
+        downsample: bool,
+        use_cls_token: bool,
+        drop_path: list[float],
+    ) -> None:
+        super().__init__()
+        self.dynamic_size = False
+        self.input_resolution = input_resolution
+        self.downsample: nn.Module
+        if downsample is True:
+            self.downsample = DeformablePatchMerging(in_dim, out_dim)
+        else:
+            self.downsample = IdentityDownsample()
+
+        blocks: list[nn.Module] = []
+        for i in range(depth):
+            if has_msa is True:
+                blocks.append(ViTBlock(out_dim, num_heads, mlp_ratio, drop_path[i]))
+            else:
+                blocks.append(MLPBlock(out_dim, mlp_ratio, drop_path[i]))
+
+        self.blocks = nn.ModuleList(blocks)
+
+        num_tokens = input_resolution[0] * input_resolution[1]
+        if use_cls_token is True:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, out_dim))
+            nn.init.trunc_normal_(self.cls_token, std=0.02)
+            num_tokens += 1
+        else:
+            self.cls_token = None
+
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_tokens, out_dim))
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+
+    def set_dynamic_size(self, dynamic_size: bool = True) -> None:
+        self.dynamic_size = dynamic_size
+
+    def _get_pos_embed(self, H: int, W: int) -> torch.Tensor:
+        if self.dynamic_size is False or (H == self.input_resolution[0] and W == self.input_resolution[1]):
+            return self.pos_embed
+
+        if self.cls_token is not None:
+            num_prefix_tokens = 1
+        else:
+            num_prefix_tokens = 0
+
+        return adjust_position_embedding(
+            self.pos_embed, self.input_resolution, (H, W), num_prefix_tokens=num_prefix_tokens
+        )
+
+    def forward(self, x: torch.Tensor, input_resolution: tuple[int, int]) -> tuple[torch.Tensor, int, int]:
+        (x, H, W) = self.downsample(x, input_resolution)
+
+        if self.cls_token is not None:
+            cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
+            x = torch.concat((cls_tokens, x), dim=1)
+
+        x = x + self._get_pos_embed(H, W)
+
+        for block in self.blocks:
+            x = block(x)
+
+        return (x, H, W)
+
+
+# pylint: disable=invalid-name
+class LIT_v1_Tiny(DetectorBackbone):
+    block_group_regex = r"body\.stage(\d+)\.blocks\.(\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 = 4
+        stage_dims: list[int] = self.config["stage_dims"]
+        depths: list[int] = self.config["depths"]
+        num_heads: list[int] = self.config["num_heads"]
+        mlp_ratios: list[float] = self.config["mlp_ratios"]
+        has_msa: list[bool] = self.config["has_msa"]
+        drop_path_rate: float = self.config["drop_path_rate"]
+
+        num_stages = len(depths)
+
+        self.stem = nn.Sequential(
+            nn.Conv2d(
+                self.input_channels,
+                stage_dims[0],
+                kernel_size=(patch_size, patch_size),
+                stride=(patch_size, patch_size),
+                padding=(0, 0),
+            ),
+            Permute([0, 2, 3, 1]),
+            nn.LayerNorm(stage_dims[0], eps=1e-6),
+        )
+
+        # Stochastic depth
+        dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(depths)).split(depths)]
+
+        stages: OrderedDict[str, nn.Module] = OrderedDict()
+        return_channels: list[int] = []
+        resolution = (self.size[0] // patch_size, self.size[1] // patch_size)
+
+        for i in range(num_stages):
+            if i > 0:
+                resolution = (resolution[0] // 2, resolution[1] // 2)
+
+            stage = LITStage(
+                stage_dims[i - 1] if i > 0 else stage_dims[0],
+                stage_dims[i],
+                input_resolution=resolution,
+                depth=depths[i],
+                num_heads=num_heads[i],
+                mlp_ratio=mlp_ratios[i],
+                has_msa=has_msa[i],
+                downsample=i > 0,
+                use_cls_token=i == num_stages - 1,
+                drop_path=dpr[i],
+            )
+            stages[f"stage{i + 1}"] = stage
+            return_channels.append(stage_dims[i])
+
+        self.body = nn.ModuleDict(stages)
+        self.norm = nn.LayerNorm(stage_dims[-1], eps=1e-6)
+        self.return_channels = return_channels
+        self.embedding_size = stage_dims[-1]
+        self.classifier = self.create_classifier()
+        self.patch_size = patch_size
+
+        # Weight initialization
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Linear):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.LayerNorm):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Conv2d):
+                if name.endswith("offset_conv") is True:
+                    continue
+
+                fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                fan_out //= m.groups
+                nn.init.normal_(m.weight, mean=0.0, std=math.sqrt(2.0 / fan_out))
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                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)
+        (B, H, W, C) = x.size()
+        x = x.reshape(B, H * W, C)
+
+        out = {}
+        for name, stage in self.body.items():
+            (x, H, W) = stage(x, (H, W))
+            if name in self.return_stages:
+                if stage.cls_token is not None:
+                    spatial_x = x[:, 1:]
+                else:
+                    spatial_x = x
+
+                out[name] = spatial_x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+
+        return out
+
+    def freeze_stages(self, up_to_stage: int) -> None:
+        for param in self.stem.parameters():
+            param.requires_grad = False
+
+        for idx, stage in enumerate(self.body.values()):
+            if idx >= up_to_stage:
+                break
+
+            for param in stage.parameters():
+                param.requires_grad = False
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.stem(x)
+        (B, H, W, C) = x.size()
+        x = x.reshape(B, H * W, C)
+        for stage in self.body.values():
+            (x, H, W) = stage(x, (H, W))
+
+        return x
+
+    def embedding(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.forward_features(x)
+        x = self.norm(x)
+        return x[:, 0]
+
+    def set_dynamic_size(self, dynamic_size: bool = True) -> None:
+        super().set_dynamic_size(dynamic_size)
+        for stage in self.body.values():
+            stage.set_dynamic_size(dynamic_size)
+
+    def adjust_size(self, new_size: tuple[int, int]) -> None:
+        if new_size == self.size:
+            return
+
+        super().adjust_size(new_size)
+
+        new_patches_resolution = (new_size[0] // self.patch_size, new_size[1] // self.patch_size)
+
+        (h, w) = new_patches_resolution
+        for stage in self.body.values():
+            if not isinstance(stage.downsample, IdentityDownsample):
+                h = h // 2
+                w = w // 2
+
+            out_resolution = (h, w)
+            if out_resolution == stage.input_resolution:
+                continue
+
+            if stage.cls_token is not None:
+                num_prefix_tokens = 1
+            else:
+                num_prefix_tokens = 0
+
+            with torch.no_grad():
+                pos_embed = adjust_position_embedding(
+                    stage.pos_embed,
+                    stage.input_resolution,
+                    out_resolution,
+                    num_prefix_tokens=num_prefix_tokens,
+                )
+
+            stage.input_resolution = out_resolution
+            stage.pos_embed = nn.Parameter(pos_embed)
+
+
+registry.register_model_config(
+    "lit_v1_t",
+    LIT_v1_Tiny,
+    config={
+        "stage_dims": [64, 128, 320, 512],
+        "depths": [3, 4, 6, 3],
+        "num_heads": [1, 2, 5, 8],
+        "mlp_ratios": [8.0, 8.0, 4.0, 4.0],
+        "has_msa": [False, False, True, True],
+        "drop_path_rate": 0.1,
+    },
+)