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

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)
@@ -66,12 +66,12 @@ class ClassAttentionBlock(nn.Module):
         self, dim: int, num_heads: int, mlp_ratio: float, qkv_bias: bool, proj_drop: float, drop_path: float, eta: float
     ) -> None:
         super().__init__()
-        self.norm1 = nn.LayerNorm(dim)
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
 
         self.attn = ClassAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, proj_drop=proj_drop)
 
         self.drop_path = StochasticDepth(drop_path, mode="row")
-        self.norm2 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
         self.mlp = MLP(dim, [int(dim * mlp_ratio), dim], activation_layer=nn.GELU, dropout=proj_drop)
 
         self.gamma1 = nn.Parameter(eta * torch.ones(dim))
@@ -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]
@@ -135,7 +135,7 @@ class LayerScaleBlock(nn.Module):
         init_values: float,
     ) -> None:
         super().__init__()
-        self.norm1 = nn.LayerNorm(dim)
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
         self.attn = TalkingHeadAttn(
             dim,
             num_heads=num_heads,
@@ -144,7 +144,7 @@ class LayerScaleBlock(nn.Module):
             proj_drop=proj_drop,
         )
         self.drop_path = StochasticDepth(drop_path, mode="row")
-        self.norm2 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
         self.mlp = MLP(dim, [int(dim * mlp_ratio), dim], activation_layer=nn.GELU, dropout=proj_drop)
         self.gamma_1 = nn.Parameter(init_values * torch.ones(dim))
         self.gamma_2 = nn.Parameter(init_values * torch.ones(dim))
@@ -221,7 +221,7 @@ class CaiT(BaseNet):
                 )
             )
 
-        self.norm = nn.LayerNorm(embed_dim)
+        self.norm = nn.LayerNorm(embed_dim, eps=1e-6)
 
         self.embedding_size = embed_dim
         self.classifier = self.create_classifier()

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

@@ -21,7 +21,7 @@ from birder.net.base import DetectorBackbone
 
 
 def insert_cls(x: torch.Tensor, cls_token: torch.Tensor) -> torch.Tensor:
-    cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+    cls_tokens = cls_token.expand(x.size(0), -1, -1)
     x = torch.concat((cls_tokens, x), dim=1)
 
     return x
@@ -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
@@ -170,7 +170,7 @@ class SerialBlock(nn.Module):
 
         # Conv-attention
         self.cpe = shared_cpe
-        self.norm1 = nn.LayerNorm(dim)
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
         self.factor_attn_crpe = FactorAttnConvRelPosEnc(
             dim,
             num_heads=num_heads,
@@ -181,7 +181,7 @@ class SerialBlock(nn.Module):
         self.drop_path = StochasticDepth(drop_path, mode="row")
 
         # MLP
-        self.norm2 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
         self.mlp = MLP(dim, [int(dim * mlp_ratio), dim], activation_layer=nn.GELU, dropout=proj_drop)
 
     def forward(self, x: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
@@ -213,9 +213,9 @@ class ParallelBlock(nn.Module):
         super().__init__()
 
         # Conv-attention
-        self.norm12 = nn.LayerNorm(dims[1])
-        self.norm13 = nn.LayerNorm(dims[2])
-        self.norm14 = nn.LayerNorm(dims[3])
+        self.norm12 = nn.LayerNorm(dims[1], eps=1e-6)
+        self.norm13 = nn.LayerNorm(dims[2], eps=1e-6)
+        self.norm14 = nn.LayerNorm(dims[3], eps=1e-6)
         self.factor_attn_crpe2 = FactorAttnConvRelPosEnc(
             dims[1], num_heads=num_heads, qkv_bias=qkv_bias, proj_drop=proj_drop, shared_crpe=shared_crpes[1]
         )
@@ -228,9 +228,9 @@ class ParallelBlock(nn.Module):
         self.drop_path = StochasticDepth(drop_path, mode="row")
 
         # MLP
-        self.norm22 = nn.LayerNorm(dims[1])
-        self.norm23 = nn.LayerNorm(dims[2])
-        self.norm24 = nn.LayerNorm(dims[3])
+        self.norm22 = nn.LayerNorm(dims[1], eps=1e-6)
+        self.norm23 = nn.LayerNorm(dims[2], eps=1e-6)
+        self.norm24 = nn.LayerNorm(dims[3], eps=1e-6)
 
         # In the parallel block, we assume dimensions are the same and share the linear transformation
         assert dims[1] == dims[2] == dims[3]
@@ -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)
@@ -447,13 +447,13 @@ class CoaT(DetectorBackbone):
 
         # Norms
         if self.parallel_blocks is not None:
-            self.norm2 = nn.LayerNorm(embed_dims[1])
-            self.norm3 = nn.LayerNorm(embed_dims[2])
+            self.norm2 = nn.LayerNorm(embed_dims[1], eps=1e-6)
+            self.norm3 = nn.LayerNorm(embed_dims[2], eps=1e-6)
         else:
             self.norm2 = None
             self.norm3 = None
 
-        self.norm4 = nn.LayerNorm(embed_dims[3])
+        self.norm4 = nn.LayerNorm(embed_dims[3], eps=1e-6)
 
         # Head
         if parallel_depth > 0:
@@ -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/conv2former.py

@@ -64,7 +64,7 @@ class SpatialAttention(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,
             ),
         )
@@ -87,8 +87,8 @@ class Conv2FormerBlock(nn.Module):
         self.mlp = MLP(dim, mlp_ratio)
 
         layer_scale_init_value = 1e-6
-        self.layer_scale_1 = nn.Parameter(layer_scale_init_value * torch.ones((1, dim, 1, 1)), requires_grad=True)
-        self.layer_scale_2 = nn.Parameter(layer_scale_init_value * torch.ones((1, dim, 1, 1)), requires_grad=True)
+        self.layer_scale_1 = nn.Parameter(layer_scale_init_value * torch.ones((1, dim, 1, 1)))
+        self.layer_scale_2 = nn.Parameter(layer_scale_init_value * torch.ones((1, dim, 1, 1)))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = x + self.drop_path(self.layer_scale_1 * self.attn(x))

birder/net/convmixer.py

@@ -58,7 +58,7 @@ class ConvMixer(BaseNet):
             inplace=None,
         )
 
-        padding = (kernel_size[0] // 2, kernel_size[1] // 2)
+        padding = ((kernel_size[0] - 1) // 2, (kernel_size[1] - 1) // 2)
         self.body = nn.Sequential(
             *[
                 nn.Sequential(

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
@@ -53,7 +45,7 @@ class ConvNeXtBlock(nn.Module):
             nn.Linear(4 * channels, channels),  # Same as 1x1 conv
             Permute([0, 3, 1, 2]),
         )
-        self.layer_scale = nn.Parameter(torch.ones(channels, 1, 1) * layer_scale, requires_grad=True)
+        self.layer_scale = nn.Parameter(torch.ones(channels, 1, 1) * layer_scale)
         self.stochastic_depth = StochasticDepth(stochastic_depth_prob, mode="row")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -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},
+)

birder/net/convnext_v2.py

@@ -56,15 +56,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
@@ -137,7 +129,7 @@ class ConvNeXt_v2(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/crossformer.py

@@ -120,9 +120,9 @@ class Attention(nn.Module):
         self.relative_position_index = nn.Buffer(relative_position_index)
 
     def forward(self, x: 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, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        (q, k, v) = qkv.unbind(0)
+        q, k, v = qkv.unbind(0)
 
         q = q * self.scale
         attn = q @ k.transpose(-2, -1)
@@ -188,15 +188,15 @@ class CrossFormerBlock(nn.Module):
         self.drop_path = StochasticDepth(drop_path, mode="row")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (H, W) = self.input_resolution
-        (B, _, C) = x.size()
+        H, W = self.input_resolution
+        B, _, C = x.size()
 
         shortcut = x
         x = self.norm1(x)
         x = x.view(B, H, W, C)
 
         # Group embeddings
-        (GH, GW) = self.group_size  # pylint: disable=invalid-name
+        GH, GW = self.group_size  # pylint: disable=invalid-name
         if self.use_lda is False:
             x = x.reshape(B, H // GH, GH, W // GW, GW, C).permute(0, 1, 3, 2, 4, 5)
         else:
@@ -244,8 +244,8 @@ class PatchMerging(nn.Module):
             )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (H, W) = self.input_resolution
-        (B, _, C) = x.shape
+        H, W = self.input_resolution
+        B, _, C = x.shape
 
         x = self.norm(x)
         x = x.view(B, H, W, C).permute(0, 3, 1, 2)
@@ -396,8 +396,8 @@ class CrossFormer(DetectorBackbone):
         for name, module in self.body.named_children():
             x = module(x)
             if name in self.return_stages:
-                (H, W) = module.resolution
-                (B, _, C) = x.size()
+                H, W = module.resolution
+                B, _, C = x.size()
                 out[name] = x.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
 
         return out

birder/net/crossvit.py

@@ -74,7 +74,7 @@ class CrossAttention(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
         # B1C -> B1H(C/H) -> BH1(C/H)
         q = self.wq(x[:, 0:1, ...]).reshape(B, 1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
         # BNC -> BNH(C/H) -> BHN(C/H)
@@ -97,7 +97,7 @@ class CrossAttentionBlock(nn.Module):
         self, dim: int, num_heads: int, qkv_bias: bool, proj_drop: float, attn_drop: float, drop_path: float
     ) -> None:
         super().__init__()
-        self.norm1 = nn.LayerNorm(dim)
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
         self.attn = CrossAttention(
             dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=proj_drop
         )
@@ -146,7 +146,7 @@ class MultiScaleBlock(nn.Module):
         for d in range(num_branches):
             self.projs.append(
                 nn.Sequential(
-                    nn.LayerNorm(dim[d]),
+                    nn.LayerNorm(dim[d], eps=1e-6),
                     nn.GELU(),
                     nn.Linear(dim[d], dim[(d + 1) % num_branches]),
                 )
@@ -187,7 +187,7 @@ class MultiScaleBlock(nn.Module):
         for d in range(num_branches):
             self.revert_projs.append(
                 nn.Sequential(
-                    nn.LayerNorm(dim[(d + 1) % num_branches]),
+                    nn.LayerNorm(dim[(d + 1) % num_branches], eps=1e-6),
                     nn.GELU(),
                     nn.Linear(dim[(d + 1) % num_branches], dim[d]),
                 )
@@ -290,7 +290,7 @@ class CrossViT(BaseNet):
             dpr_ptr += curr_depth
             self.blocks.append(block)
 
-        self.norm = nn.ModuleList([nn.LayerNorm(embed_dim[i]) for i in range(self.num_branches)])
+        self.norm = nn.ModuleList([nn.LayerNorm(embed_dim[i], eps=1e-6) for i in range(self.num_branches)])
         self.embedding_size = sum(self.embed_dim)
         self.classifier = nn.ModuleList()
         for i in range(self.num_branches):
@@ -482,7 +482,7 @@ registry.register_weights(
         "formats": {
             "pt": {
                 "file_size": 32.7,
-                "sha256": "515265ed725adce09464bfd23ce612b1d1178bc22a57960db089d7148556149a",
+                "sha256": "08f674d8165dc97cc535f8188a5c5361751a8d0bb85061454986a21541a6fe8e",
             }
         },
         "net": {"network": "crossvit_9d", "tag": "il-common"},

birder/net/cspnet.py

@@ -226,7 +226,7 @@ class CrossStage(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.conv_down(x)
         x = self.conv_exp(x)
-        (xs, xb) = x.split(self.expand_channels // 2, dim=1)
+        xs, xb = x.split(self.expand_channels // 2, dim=1)
         xb = self.blocks(xb)
         xb = self.conv_transition_b(xb).contiguous()
         out = self.conv_transition(torch.concat([xs, xb], dim=1))