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

birder/net/_vit_configs.py

@@ -0,0 +1,479 @@
+"""
+ViT model configuration registrations
+
+This file contains *only* model variant definitions and their registration
+with the global model registry. The actual ViT implementation lives in vit.py.
+
+Naming:
+- All model names must follow the ViT / RoPE ViT naming convention documented in rope_vit_configs.py.
+"""
+
+from birder.model_registry import registry
+from birder.net.base import BaseNet
+
+TINY = {"num_layers": 12, "num_heads": 3, "hidden_dim": 192, "mlp_dim": 768, "drop_path_rate": 0.0}
+SMALL = {"num_layers": 12, "num_heads": 6, "hidden_dim": 384, "mlp_dim": 1536, "drop_path_rate": 0.0}
+MEDIUM = {"num_layers": 12, "num_heads": 8, "hidden_dim": 512, "mlp_dim": 2048, "drop_path_rate": 0.0}
+BASE = {"num_layers": 12, "num_heads": 12, "hidden_dim": 768, "mlp_dim": 3072, "drop_path_rate": 0.1}
+LARGE = {"num_layers": 24, "num_heads": 16, "hidden_dim": 1024, "mlp_dim": 4096, "drop_path_rate": 0.1}
+HUGE = {"num_layers": 32, "num_heads": 16, "hidden_dim": 1280, "mlp_dim": 5120, "drop_path_rate": 0.1}
+
+# From "Getting vit in Shape: Scaling Laws for Compute-Optimal Model Design"
+# Shape-optimized vision transformer (SoViT)
+SO150 = {
+    "num_layers": 18,
+    "num_heads": 16,
+    "hidden_dim": 896,  # Changed from 880 for RoPE divisibility
+    "mlp_dim": 2320,
+    "drop_path_rate": 0.1,
+}
+SO400 = {
+    "num_layers": 27,
+    "num_heads": 16,
+    "hidden_dim": 1152,
+    "mlp_dim": 4304,
+    "drop_path_rate": 0.1,
+}
+
+# From "Scaling Vision Transformers"
+GIANT = {"num_layers": 40, "num_heads": 16, "hidden_dim": 1408, "mlp_dim": 6144, "drop_path_rate": 0.1}
+GIGANTIC = {"num_layers": 48, "num_heads": 16, "hidden_dim": 1664, "mlp_dim": 8192, "drop_path_rate": 0.1}
+
+
+def register_vit_configs(vit: type[BaseNet]) -> None:
+    registry.register_model_config(
+        "vit_t32",
+        vit,
+        config={"patch_size": 32, **TINY},
+    )
+    registry.register_model_config(
+        "vit_t16",
+        vit,
+        config={"patch_size": 16, **TINY},
+    )
+    registry.register_model_config(
+        "vit_t14",
+        vit,
+        config={"patch_size": 14, **TINY},
+    )
+    registry.register_model_config(
+        "vit_s32",
+        vit,
+        config={"patch_size": 32, **SMALL},
+    )
+    registry.register_model_config(
+        "vit_s16",
+        vit,
+        config={"patch_size": 16, **SMALL},
+    )
+    registry.register_model_config(
+        "vit_s16_ls",
+        vit,
+        config={"patch_size": 16, **SMALL, "layer_scale_init_value": 1e-5},
+    )
+    registry.register_model_config(
+        "vit_s16_pn",
+        vit,
+        config={"patch_size": 16, **SMALL, "pre_norm": True, "norm_layer_eps": 1e-5},
+    )
+    registry.register_model_config(
+        "vit_s14",
+        vit,
+        config={"patch_size": 14, **SMALL},
+    )
+    registry.register_model_config(
+        "vit_m32",
+        vit,
+        config={"patch_size": 32, **MEDIUM},
+    )
+    registry.register_model_config(
+        "vit_m16",
+        vit,
+        config={"patch_size": 16, **MEDIUM},
+    )
+    registry.register_model_config(
+        "vit_m14",
+        vit,
+        config={"patch_size": 14, **MEDIUM},
+    )
+    registry.register_model_config(
+        "vit_b32",
+        vit,
+        config={"patch_size": 32, **BASE, "drop_path_rate": 0.0},  # Override the BASE definition
+    )
+    registry.register_model_config(
+        "vit_b16",
+        vit,
+        config={"patch_size": 16, **BASE},
+    )
+    registry.register_model_config(
+        "vit_b16_ls",
+        vit,
+        config={"patch_size": 16, **BASE, "layer_scale_init_value": 1e-5},
+    )
+    registry.register_model_config(
+        "vit_b16_qkn_ls",
+        vit,
+        config={"patch_size": 16, **BASE, "layer_scale_init_value": 1e-5, "qk_norm": True},
+    )
+    registry.register_model_config(
+        "vit_b16_pn_quick_gelu",
+        vit,
+        config={"patch_size": 16, **BASE, "pre_norm": True, "norm_layer_eps": 1e-5, "act_layer_type": "quick_gelu"},
+    )
+    registry.register_model_config(
+        "vit_b14",
+        vit,
+        config={"patch_size": 14, **BASE},
+    )
+    registry.register_model_config(
+        "vit_so150m_p14_avg",
+        vit,
+        config={"patch_size": 14, **SO150, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_so150m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO150, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_l32",
+        vit,
+        config={"patch_size": 32, **LARGE},
+    )
+    registry.register_model_config(
+        "vit_l16",
+        vit,
+        config={"patch_size": 16, **LARGE},
+    )
+    registry.register_model_config(
+        "vit_l14",
+        vit,
+        config={"patch_size": 14, **LARGE},
+    )
+    registry.register_model_config(
+        "vit_l14_pn",
+        vit,
+        config={"patch_size": 14, **LARGE, "pre_norm": True, "norm_layer_eps": 1e-5},
+    )
+    registry.register_model_config(
+        "vit_l14_pn_quick_gelu",
+        vit,
+        config={"patch_size": 14, **LARGE, "pre_norm": True, "norm_layer_eps": 1e-5, "act_layer_type": "quick_gelu"},
+    )
+    registry.register_model_config(
+        "vit_so400m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO400, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_h16",
+        vit,
+        config={"patch_size": 16, **HUGE},
+    )
+    registry.register_model_config(
+        "vit_h14",
+        vit,
+        config={"patch_size": 14, **HUGE},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_g16",
+        vit,
+        config={"patch_size": 16, **GIANT},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_g14",
+        vit,
+        config={"patch_size": 14, **GIANT},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_gigantic14",
+        vit,
+        config={"patch_size": 14, **GIGANTIC},
+    )
+    registry.register_model_config(  # From "PaLI: A Jointly-Scaled Multilingual Language-Image Model"
+        "vit_e14",
+        vit,
+        config={
+            "patch_size": 14,
+            "num_layers": 56,
+            "num_heads": 16,
+            "hidden_dim": 1792,
+            "mlp_dim": 15360,
+            "drop_path_rate": 0.1,
+        },
+    )
+    registry.register_model_config(  # From "Scaling Language-Free Visual Representation Learning"
+        "vit_1b_p16",  # AKA vit_giant2 from DINOv2
+        vit,
+        config={
+            "patch_size": 16,
+            "num_layers": 40,
+            "num_heads": 24,
+            "hidden_dim": 1536,
+            "mlp_dim": 6144,
+            "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
+    ####################
+
+    registry.register_model_config(
+        "vit_reg1_t32",
+        vit,
+        config={"patch_size": 32, **TINY, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_t16",
+        vit,
+        config={"patch_size": 16, **TINY, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_t14",
+        vit,
+        config={"patch_size": 14, **TINY, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_s32",
+        vit,
+        config={"patch_size": 32, **SMALL, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_s16",
+        vit,
+        config={"patch_size": 16, **SMALL, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_s16_ls",
+        vit,
+        config={"patch_size": 16, **SMALL, "layer_scale_init_value": 1e-5, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_s16_rms_ls",
+        vit,
+        config={
+            "patch_size": 16,
+            **SMALL,
+            "layer_scale_init_value": 1e-5,
+            "num_reg_tokens": 1,
+            "norm_layer_type": "RMSNorm",
+        },
+    )
+    registry.register_model_config(
+        "vit_reg1_s14",
+        vit,
+        config={"patch_size": 14, **SMALL, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg4_m32",
+        vit,
+        config={"patch_size": 32, **MEDIUM, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_m16",
+        vit,
+        config={"patch_size": 16, **MEDIUM, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_m16_rms_avg",
+        vit,
+        config={"patch_size": 16, **MEDIUM, "num_reg_tokens": 4, "class_token": False, "norm_layer_type": "RMSNorm"},
+    )
+    registry.register_model_config(
+        "vit_reg4_m14",
+        vit,
+        config={"patch_size": 14, **MEDIUM, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_b32",
+        vit,
+        config={"patch_size": 32, **BASE, "num_reg_tokens": 4, "drop_path_rate": 0.0},  # Override the BASE definition
+    )
+    registry.register_model_config(
+        "vit_reg4_b16",
+        vit,
+        config={"patch_size": 16, **BASE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_b16_avg",
+        vit,
+        config={"patch_size": 16, **BASE, "num_reg_tokens": 4, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg4_b14",
+        vit,
+        config={"patch_size": 14, **BASE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg8_b14_ap",
+        vit,
+        config={"patch_size": 14, **BASE, "num_reg_tokens": 8, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p16_avg",
+        vit,
+        config={"patch_size": 16, **SO150, "num_reg_tokens": 4, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p16_swiglu_ap",
+        vit,
+        config={
+            "patch_size": 16,
+            **SO150,
+            "num_reg_tokens": 8,
+            "class_token": False,
+            "attn_pool_head": True,
+            "mlp_layer_type": "SwiGLU_FFN",
+        },
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p14_avg",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 4, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p14_ls",
+        vit,
+        config={"patch_size": 14, **SO150, "layer_scale_init_value": 1e-5, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 4, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p14_aps",
+        vit,
+        config={
+            "patch_size": 14,
+            **SO150,
+            "num_reg_tokens": 4,
+            "class_token": False,
+            "attn_pool_head": True,
+            "attn_pool_special_tokens": True,
+        },
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p14_avg",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 8, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p14_swiglu",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 8, "mlp_layer_type": "SwiGLU_FFN"},
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p14_swiglu_avg",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 8, "class_token": False, "mlp_layer_type": "SwiGLU_FFN"},
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 8, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg4_l32",
+        vit,
+        config={"patch_size": 32, **LARGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_l16",
+        vit,
+        config={"patch_size": 16, **LARGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg8_l16_avg",
+        vit,
+        config={"patch_size": 16, **LARGE, "num_reg_tokens": 8, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg8_l16_aps",
+        vit,
+        config={
+            "patch_size": 16,
+            **LARGE,
+            "num_reg_tokens": 8,
+            "class_token": False,
+            "attn_pool_head": True,
+            "attn_pool_special_tokens": True,
+        },
+    )
+    registry.register_model_config(
+        "vit_reg4_l14",
+        vit,
+        config={"patch_size": 14, **LARGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(  # DeiT III style
+        "vit_reg4_l14_nps_ls",
+        vit,
+        config={
+            "pos_embed_special_tokens": False,
+            "patch_size": 14,
+            **LARGE,
+            "layer_scale_init_value": 1e-5,
+            "num_reg_tokens": 4,
+        },
+    )
+    registry.register_model_config(
+        "vit_reg8_l14_ap",
+        vit,
+        config={"patch_size": 14, **LARGE, "num_reg_tokens": 8, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg8_l14_rms_ap",
+        vit,
+        config={
+            "patch_size": 14,
+            **LARGE,
+            "num_reg_tokens": 8,
+            "class_token": False,
+            "attn_pool_head": True,
+            "norm_layer_type": "RMSNorm",
+        },
+    )
+    registry.register_model_config(
+        "vit_reg8_so400m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO400, "num_reg_tokens": 8, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg4_h16",
+        vit,
+        config={"patch_size": 16, **HUGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_h14",
+        vit,
+        config={"patch_size": 14, **HUGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_reg4_g16",
+        vit,
+        config={"patch_size": 16, **GIANT, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_reg4_g14",
+        vit,
+        config={"patch_size": 14, **GIANT, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_reg4_gigantic14",
+        vit,
+        config={"patch_size": 14, **GIGANTIC, "num_reg_tokens": 4},
+    )

birder/net/biformer.py

@@ -8,6 +8,7 @@ Changes from original:
 * All attention types are in (B, C, H, W)
 * Using the newer Bi-Level Routing Attention implementation
 * Dynamic n_win size (image size // 32)
+* Stem bias term removed
 """
 
 # Reference license: Apache-2.0

birder/net/cait.py

@@ -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))
@@ -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/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
@@ -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]
@@ -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:

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

@@ -53,7 +53,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:

birder/net/crossvit.py

@@ -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/davit.py

@@ -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: