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

birder/net/simple_vit.py

@@ -20,6 +20,12 @@ import torch
 from torch import nn
 
 from birder.model_registry import registry
+from birder.net._vit_configs import BASE
+from birder.net._vit_configs import GIANT
+from birder.net._vit_configs import HUGE
+from birder.net._vit_configs import LARGE
+from birder.net._vit_configs import MEDIUM
+from birder.net._vit_configs import SMALL
 from birder.net.base import MaskedTokenOmissionMixin
 from birder.net.base import PreTrainEncoder
 from birder.net.base import TokenOmissionResultType
@@ -45,12 +51,12 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
         assert self.config is not None, "must set config"
 
         image_size = self.size
-        drop_path_rate = 0.0
         patch_size: int = self.config["patch_size"]
         num_layers: int = self.config["num_layers"]
         num_heads: int = self.config["num_heads"]
         hidden_dim: int = self.config["hidden_dim"]
         mlp_dim: int = self.config["mlp_dim"]
+        drop_path_rate: float = self.config["drop_path_rate"]
 
         torch._assert(image_size[0] % patch_size == 0, "Input shape indivisible by patch size!")
         torch._assert(image_size[1] % patch_size == 0, "Input shape indivisible by patch size!")
@@ -215,75 +221,75 @@ class Simple_ViT(PreTrainEncoder, MaskedTokenOmissionMixin):
 registry.register_model_config(
     "simple_vit_s32",
     Simple_ViT,
-    config={"patch_size": 32, "num_layers": 12, "num_heads": 6, "hidden_dim": 384, "mlp_dim": 1536},
+    config={"patch_size": 32, **SMALL},
 )
 registry.register_model_config(
     "simple_vit_s16",
     Simple_ViT,
-    config={"patch_size": 16, "num_layers": 12, "num_heads": 6, "hidden_dim": 384, "mlp_dim": 1536},
+    config={"patch_size": 16, **SMALL},
 )
 registry.register_model_config(
     "simple_vit_s14",
     Simple_ViT,
-    config={"patch_size": 14, "num_layers": 12, "num_heads": 6, "hidden_dim": 384, "mlp_dim": 1536},
+    config={"patch_size": 14, **SMALL},
 )
 registry.register_model_config(
     "simple_vit_m32",
     Simple_ViT,
-    config={"patch_size": 32, "num_layers": 12, "num_heads": 8, "hidden_dim": 512, "mlp_dim": 2048},
+    config={"patch_size": 32, **MEDIUM},
 )
 registry.register_model_config(
     "simple_vit_m16",
     Simple_ViT,
-    config={"patch_size": 16, "num_layers": 12, "num_heads": 8, "hidden_dim": 512, "mlp_dim": 2048},
+    config={"patch_size": 16, **MEDIUM},
 )
 registry.register_model_config(
     "simple_vit_m14",
     Simple_ViT,
-    config={"patch_size": 14, "num_layers": 12, "num_heads": 8, "hidden_dim": 512, "mlp_dim": 2048},
+    config={"patch_size": 14, **MEDIUM},
 )
 registry.register_model_config(
     "simple_vit_b32",
     Simple_ViT,
-    config={"patch_size": 32, "num_layers": 12, "num_heads": 12, "hidden_dim": 768, "mlp_dim": 3072},
+    config={"patch_size": 32, **BASE},  # Override the BASE definition
 )
 registry.register_model_config(
     "simple_vit_b16",
     Simple_ViT,
-    config={"patch_size": 16, "num_layers": 12, "num_heads": 12, "hidden_dim": 768, "mlp_dim": 3072},
+    config={"patch_size": 16, **BASE},
 )
 registry.register_model_config(
     "simple_vit_b14",
     Simple_ViT,
-    config={"patch_size": 14, "num_layers": 12, "num_heads": 12, "hidden_dim": 768, "mlp_dim": 3072},
+    config={"patch_size": 14, **BASE},
 )
 registry.register_model_config(
     "simple_vit_l32",
     Simple_ViT,
-    config={"patch_size": 32, "num_layers": 24, "num_heads": 16, "hidden_dim": 1024, "mlp_dim": 4096},
+    config={"patch_size": 32, **LARGE},
 )
 registry.register_model_config(
     "simple_vit_l16",
     Simple_ViT,
-    config={"patch_size": 16, "num_layers": 24, "num_heads": 16, "hidden_dim": 1024, "mlp_dim": 4096},
+    config={"patch_size": 16, **LARGE},
 )
 registry.register_model_config(
     "simple_vit_l14",
     Simple_ViT,
-    config={"patch_size": 14, "num_layers": 24, "num_heads": 16, "hidden_dim": 1024, "mlp_dim": 4096},
+    config={"patch_size": 14, **LARGE},
 )
 registry.register_model_config(
     "simple_vit_h16",
     Simple_ViT,
-    config={"patch_size": 16, "num_layers": 32, "num_heads": 16, "hidden_dim": 1280, "mlp_dim": 5120},
+    config={"patch_size": 16, **HUGE},
 )
 registry.register_model_config(
     "simple_vit_h14",
     Simple_ViT,
-    config={"patch_size": 14, "num_layers": 32, "num_heads": 16, "hidden_dim": 1280, "mlp_dim": 5120},
+    config={"patch_size": 14, **HUGE},
 )
 registry.register_model_config(  # From "Scaling Vision Transformers"
     "simple_vit_g14",
     Simple_ViT,
-    config={"patch_size": 14, "num_layers": 40, "num_heads": 16, "hidden_dim": 1408, "mlp_dim": 6144},
+    config={"patch_size": 14, **GIANT},
 )

birder/net/smt.py

@@ -259,7 +259,7 @@ class Stem(nn.Module):
                 embed_dim,
                 kernel_size=kernel_size,
                 stride=stride,
-                padding=(kernel_size[0] // 2, kernel_size[1] // 2),
+                padding=((kernel_size[0] - 1) // 2, (kernel_size[1] - 1) // 2),
             ),
             nn.Conv2d(embed_dim, embed_dim, kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)),
         )

birder/net/squeezenet.py

@@ -20,11 +20,11 @@ from birder.net.base import BaseNet
 class Fire(nn.Module):
     def __init__(self, in_planes: int, squeeze: int, expand: int) -> None:
         super().__init__()
-        self.squeeze = nn.Conv2d(in_planes, squeeze, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=False)
+        self.squeeze = nn.Conv2d(in_planes, squeeze, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
         self.squeeze_activation = nn.ReLU(inplace=True)
-        self.left = nn.Conv2d(squeeze, expand, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=False)
+        self.left = nn.Conv2d(squeeze, expand, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
         self.left_activation = nn.ReLU(inplace=True)
-        self.right = nn.Conv2d(squeeze, expand, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+        self.right = nn.Conv2d(squeeze, expand, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
         self.right_activation = nn.ReLU(inplace=True)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -53,7 +53,7 @@ class SqueezeNet(BaseNet):
         assert self.config is None, "config not supported"
 
         self.stem = nn.Sequential(
-            nn.Conv2d(self.input_channels, 64, kernel_size=(3, 3), stride=(2, 2), padding=(0, 0), bias=False),
+            nn.Conv2d(self.input_channels, 64, kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(0, 0), ceil_mode=True),
         )
@@ -94,14 +94,7 @@ class SqueezeNet(BaseNet):
 
         return nn.Sequential(
             nn.Dropout(p=0.5, inplace=True),
-            nn.Conv2d(
-                embed_dim,
-                self.num_classes,
-                kernel_size=(1, 1),
-                stride=(1, 1),
-                padding=(0, 0),
-                bias=False,
-            ),
+            nn.Conv2d(embed_dim, self.num_classes, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0)),
             nn.ReLU(inplace=True),
             nn.AdaptiveAvgPool2d(output_size=(1, 1)),
             nn.Flatten(1),

birder/net/squeezenext.py

@@ -27,7 +27,6 @@ class SqnxtUnit(nn.Module):
                 kernel_size=(1, 1),
                 stride=(stride, stride),
                 padding=(0, 0),
-                bias=False,
             )
 
         elif in_channels > out_channels:
@@ -38,7 +37,6 @@ class SqnxtUnit(nn.Module):
                 kernel_size=(1, 1),
                 stride=(stride, stride),
                 padding=(0, 0),
-                bias=False,
             )
 
         else:
@@ -52,7 +50,6 @@ class SqnxtUnit(nn.Module):
                 kernel_size=(1, 1),
                 stride=(stride, stride),
                 padding=(0, 0),
-                bias=False,
             ),
             Conv2dNormActivation(
                 in_channels // reduction,
@@ -60,7 +57,6 @@ class SqnxtUnit(nn.Module):
                 kernel_size=(1, 1),
                 stride=(1, 1),
                 padding=(0, 0),
-                bias=False,
             ),
             Conv2dNormActivation(
                 in_channels // (2 * reduction),
@@ -68,7 +64,6 @@ class SqnxtUnit(nn.Module):
                 kernel_size=(1, 3),
                 stride=(1, 1),
                 padding=(0, 1),
-                bias=False,
             ),
             Conv2dNormActivation(
                 in_channels // reduction,
@@ -76,7 +71,6 @@ class SqnxtUnit(nn.Module):
                 kernel_size=(3, 1),
                 stride=(1, 1),
                 padding=(1, 0),
-                bias=False,
             ),
             Conv2dNormActivation(
                 in_channels // reduction,
@@ -84,7 +78,6 @@ class SqnxtUnit(nn.Module):
                 kernel_size=(1, 1),
                 stride=(1, 1),
                 padding=(0, 0),
-                bias=False,
             ),
         )
         self.relu = nn.ReLU(inplace=True)
@@ -124,7 +117,6 @@ class SqueezeNext(DetectorBackbone):
                 kernel_size=(7, 7),
                 stride=(2, 2),
                 padding=(1, 1),
-                bias=False,
             ),
             nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(0, 0), ceil_mode=True),
         )
@@ -155,7 +147,6 @@ class SqueezeNext(DetectorBackbone):
                 kernel_size=(1, 1),
                 stride=(1, 1),
                 padding=(0, 0),
-                bias=False,
             ),
             nn.AdaptiveAvgPool2d(output_size=(1, 1)),
             nn.Flatten(1),
@@ -199,18 +190,3 @@ registry.register_model_config("squeezenext_0_5", SqueezeNext, config={"width_sc
 registry.register_model_config("squeezenext_1_0", SqueezeNext, config={"width_scale": 1.0})
 registry.register_model_config("squeezenext_1_5", SqueezeNext, config={"width_scale": 1.5})
 registry.register_model_config("squeezenext_2_0", SqueezeNext, config={"width_scale": 2.0})
-
-registry.register_weights(
-    "squeezenext_1_0_il-common",
-    {
-        "description": "SqueezeNext v2 1.0x output channels model trained on the il-common dataset",
-        "resolution": (259, 259),
-        "formats": {
-            "pt": {
-                "file_size": 3.5,
-                "sha256": "da01d1cd05c71b80b5e4e6ca66400f64fa3f6179d0e90834c4f6942c8095557a",
-            }
-        },
-        "net": {"network": "squeezenext_1_0", "tag": "il-common"},
-    },
-)

birder/net/ssl/capi.py

@@ -306,7 +306,7 @@ class Decoder(nn.Module):
             dim=decoder_embed_dim,
             num_special_tokens=0,
         ).unsqueeze(0)
-        self.decoder_pos_embed = nn.Parameter(pos_embedding, requires_grad=False)
+        self.decoder_pos_embed = nn.Buffer(pos_embedding)
 
         self.decoder_layers = nn.ModuleList()
         for _ in range(decoder_depth):

birder/net/ssl/data2vec.py

@@ -51,7 +51,7 @@ class Data2Vec(SSLBaseNet):
         self.ema_backbone = copy.deepcopy(self.backbone)
         self.head = nn.Linear(self.backbone.embedding_size, self.backbone.embedding_size)
 
-        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width), requires_grad=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width))
 
         # Weights initialization
         self.ema_backbone.load_state_dict(self.backbone.state_dict())

birder/net/ssl/dino_v2.py

@@ -460,7 +460,7 @@ class DINOv2Student(SSLBaseNet):
                 bottleneck_dim=head_bottleneck_dim,
             )
 
-        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width), requires_grad=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width))
 
     # pylint: disable=arguments-differ
     def forward(  # type: ignore[override]
@@ -543,7 +543,7 @@ class DINOv2Teacher(SSLBaseNet):
             )
 
         # Unused, Makes for an easier EMA update
-        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width), requires_grad=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width))
 
     # pylint: disable=arguments-differ
     def forward(  # type: ignore[override]

birder/net/ssl/franca.py

@@ -433,7 +433,7 @@ class FrancaStudent(SSLBaseNet):
                 nesting_list=nesting_list,
             )
 
-        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width), requires_grad=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width))
 
     # pylint: disable=arguments-differ
     def forward(  # type: ignore[override]
@@ -523,7 +523,7 @@ class FrancaTeacher(SSLBaseNet):
             )
 
         # Unused, Makes for an easier EMA update
-        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width), requires_grad=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width))
 
     # pylint: disable=arguments-differ
     def forward(  # type: ignore[override]

birder/net/ssl/i_jepa.py

@@ -200,7 +200,7 @@ class VisionTransformerPredictor(nn.Module):
         self.mask_token = nn.Parameter(torch.zeros(1, 1, predictor_embed_dim))
 
         pos_embedding = pos_embedding_sin_cos_2d(h=size[0], w=size[1], dim=predictor_embed_dim, num_special_tokens=0)
-        self.pos_embedding = nn.Parameter(pos_embedding, requires_grad=False)
+        self.pos_embedding = nn.Buffer(pos_embedding)
 
         self.encoder = Encoder(
             depth, num_heads, predictor_embed_dim, mlp_dim, dropout=0.0, attention_dropout=0.0, dpr=dpr

birder/net/ssl/ibot.py

@@ -254,7 +254,7 @@ class iBOT(SSLBaseNet):
             shared_head=shared_head,
         )
 
-        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width), requires_grad=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, self.backbone.stem_width))
 
     def forward(  # type: ignore[override]  # pylint: disable=arguments-differ
         self, x: torch.Tensor, masks: Optional[torch.Tensor], return_keys: Literal["all", "embedding"] = "all"

birder/net/swiftformer.py

@@ -48,7 +48,12 @@ class ConvEncoder(nn.Module):
     ) -> None:
         super().__init__()
         self.dw_conv = nn.Conv2d(
-            dim, dim, kernel_size, stride=(1, 1), padding=(kernel_size[0] // 2, kernel_size[1] // 2), groups=dim
+            dim,
+            dim,
+            kernel_size,
+            stride=(1, 1),
+            padding=((kernel_size[0] - 1) // 2, (kernel_size[1] - 1) // 2),
+            groups=dim,
         )
         self.norm = nn.BatchNorm2d(dim)
         self.pw_conv1 = nn.Conv2d(dim, hidden_dim, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
@@ -125,7 +130,12 @@ class LocalRepresentation(nn.Module):
     def __init__(self, dim: int, kernel_size: tuple[int, int], drop_path: float, use_layer_scale: bool) -> None:
         super().__init__()
         self.dw_conv = nn.Conv2d(
-            dim, dim, kernel_size, stride=(1, 1), padding=(kernel_size[0] // 2, kernel_size[1] // 2), groups=dim
+            dim,
+            dim,
+            kernel_size,
+            stride=(1, 1),
+            padding=((kernel_size[0] - 1) // 2, (kernel_size[1] - 1) // 2),
+            groups=dim,
         )
         self.norm = nn.BatchNorm2d(dim)
         self.pw_conv1 = nn.Conv2d(dim, dim, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))

birder/net/swin_transformer_v2.py

@@ -72,7 +72,7 @@ class ShiftedWindowAttention(nn.Module):
         self.define_relative_position_bias_table()
         self.define_relative_position_index()
 
-        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)
+        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
 
         # MLP to generate continuous relative position bias
         self.cpb_mlp = nn.Sequential(

birder/net/tiny_vit.py

@@ -77,10 +77,11 @@ class MBConv(nn.Module):
             kernel_size=(1, 1),
             stride=(1, 1),
             padding=(0, 0),
-            activation_layer=nn.GELU,
+            activation_layer=None,
             inplace=None,
         )
         self.drop_path = StochasticDepth(drop_path, mode="row")
+        self.act = nn.GELU()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         shortcut = x
@@ -89,6 +90,7 @@ class MBConv(nn.Module):
         x = self.conv3(x)
         x = self.drop_path(x)
         x += shortcut
+        x = self.act(x)
 
         return x
 
@@ -508,18 +510,3 @@ registry.register_model_config(
         "drop_path_rate": 0.2,
     },
 )
-
-registry.register_weights(
-    "tiny_vit_5m_il-common",
-    {
-        "description": "TinyViT 5M model trained on the il-common dataset",
-        "resolution": (256, 256),
-        "formats": {
-            "pt": {
-                "file_size": 20.0,
-                "sha256": "57f84dc3144fc4e3ca39328d3a1446ca9e26ddb54e4c4d84301b7638bee2ec21",
-            },
-        },
-        "net": {"network": "tiny_vit_5m", "tag": "il-common"},
-    },
-)

birder/net/van.py

@@ -116,8 +116,8 @@ class VANBlock(nn.Module):
         self.mlp = DWConvMLP(in_features=dim, hidden_features=int(dim * mlp_ratio), drop=drop)
 
         layer_scale_init_value = 1e-2
-        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(self.norm1(x)))