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

birder/net/smt.py

@@ -36,7 +36,7 @@ class DWConv(nn.Module):
         self.dwconv = nn.Conv2d(dim, dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=dim)
 
     def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
-        (B, _, C) = x.size()
+        B, _, C = x.size()
         x = x.transpose(1, 2).view(B, C, H, W)
         x = self.dwconv(x)
         x = x.flatten(2).transpose(1, 2)
@@ -94,7 +94,7 @@ class CAAttention(nn.Module):
         self.proj1 = nn.Conv2d(dim * expand_ratio, dim, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
 
     def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
-        (B, N, C) = x.size()
+        B, N, C = x.size()
 
         v = self.v(x)
         s = self.s(x).reshape(B, H, W, self.ca_num_heads, C // self.ca_num_heads).permute(3, 0, 4, 1, 2)
@@ -140,11 +140,11 @@ class SAAttention(nn.Module):
         self.conv = nn.Conv2d(dim, dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=dim)
 
     def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
-        (B, N, C) = x.size()
+        B, N, C = x.size()
 
         q = self.q(x).reshape(B, N, self.sa_num_heads, C // self.sa_num_heads).permute(0, 2, 1, 3)
         kv = self.kv(x).reshape(B, -1, 2, self.sa_num_heads, C // self.sa_num_heads).permute(2, 0, 3, 1, 4)
-        (k, v) = kv.unbind(0)
+        k, v = kv.unbind(0)
         attn = (q @ k.transpose(-2, -1)) * self.scale
         attn = attn.softmax(dim=-1)
         attn = self.attn_drop(attn)
@@ -243,7 +243,7 @@ class OverlapPatchEmbed(nn.Module):
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int]:
         x = self.proj(x)
-        (_, _, H, W) = x.size()
+        _, _, H, W = x.size()
         x = x.flatten(2).transpose(1, 2)
         x = self.norm(x)
 
@@ -267,7 +267,7 @@ class Stem(nn.Module):
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int]:
         x = self.conv(x)
-        (_, _, H, W) = x.size()
+        _, _, H, W = x.size()
         x = x.flatten(2).transpose(1, 2)
         x = self.norm(x)
 
@@ -329,7 +329,7 @@ class SMTStage(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         B = x.size(0)
-        (x, H, W) = self.downsample_block(x)
+        x, H, W = self.downsample_block(x)
         x = self.blocks(x, H, W)
         x = self.norm(x)
         x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()

birder/net/ssl/barlow_twins.py

@@ -21,7 +21,7 @@ from birder.net.ssl.base import SSLBaseNet
 
 def off_diagonal(x: torch.Tensor) -> torch.Tensor:
     # Return a flattened view of the off-diagonal elements of a square matrix
-    (n, _) = x.size()
+    n, _ = x.size()
     # assert n == m
     return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
 

birder/net/ssl/byol.py

@@ -80,11 +80,11 @@ class BYOL(SSLBaseNet):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         projection = self.online_encoder(x)
         online_predictions = self.online_predictor(projection)
-        (online_pred_one, online_pred_two) = online_predictions.chunk(2, dim=0)
+        online_pred_one, online_pred_two = online_predictions.chunk(2, dim=0)
 
         with torch.no_grad():
             target_projections = self.target_encoder(x)
-            (target_proj_one, target_proj_two) = target_projections.chunk(2, dim=0)
+            target_proj_one, target_proj_two = target_projections.chunk(2, dim=0)
 
         loss_one = loss_fn(online_pred_one, target_proj_two.detach())
         loss_two = loss_fn(online_pred_two, target_proj_one.detach())

birder/net/ssl/capi.py

@@ -263,11 +263,11 @@ class CrossAttention(nn.Module):
         self.proj = nn.Linear(decoder_dim, decoder_dim)
 
     def forward(self, tgt: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
-        (B, N, C) = tgt.size()
+        B, N, C = tgt.size()
         n_kv = memory.size(1)
         q = self.q(tgt).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
         kv = self.kv(memory).reshape(B, n_kv, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        (k, v) = kv.unbind(0)
+        k, v = kv.unbind(0)
 
         attn = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)  # pylint: disable=not-callable
         x = attn.transpose(1, 2).reshape(B, N, C)
@@ -419,7 +419,7 @@ class CAPITeacher(SSLBaseNet):
             x = self.backbone.masked_encoding_omission(x, ids_keep)["tokens"]
 
         x = x[:, self.backbone.num_special_tokens :, :]
-        (assignments, clustering_loss) = self.head(x.transpose(0, 1))
+        assignments, clustering_loss = self.head(x.transpose(0, 1))
 
         assignments = assignments.detach().transpose(0, 1)
         row_indices = torch.arange(B).unsqueeze(1).expand_as(ids_predict)

birder/net/ssl/data2vec2.py

@@ -68,7 +68,7 @@ class Decoder2d(nn.Module):
         self.proj = nn.Linear(embed_dim, in_channels)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, _, C) = x.size()  # B, N, C
+        B, _, C = x.size()  # B, N, C
 
         x = x.transpose(1, 2).reshape(B, C, self.H, self.W)
 

birder/net/ssl/dino_v2.py

@@ -148,7 +148,17 @@ class DINOLoss(nn.Module):
 
     def forward(
         self, student_output_list: list[torch.Tensor], teacher_out_softmax_centered_list: list[torch.Tensor]
-    ) -> float:
+    ) -> torch.Tensor:
+        s = torch.stack(student_output_list, 0)
+        t = torch.stack(teacher_out_softmax_centered_list, 0)
+        lsm = F.log_softmax(s / self.student_temp, dim=-1)
+        loss = -(torch.einsum("tbk,sbk->tsb", t, lsm).mean(-1).sum())
+
+        return loss
+
+    def forward_reference(
+        self, student_output_list: list[torch.Tensor], teacher_out_softmax_centered_list: list[torch.Tensor]
+    ) -> torch.Tensor:
         total_loss = 0.0
         for s in student_output_list:
             lsm = F.log_softmax(s / self.student_temp, dim=-1)

birder/net/ssl/franca.py

@@ -69,7 +69,7 @@ class DINOHeadMRL(nn.Module):
     ) -> None:
         super().__init__()
         self.nesting_list = nesting_list
-        self.matryoshka_projections = nn.ModuleList([nn.Linear(dim, dim, bias=True) for dim in self.nesting_list])
+        self.matryoshka_projections = nn.ModuleList([nn.Linear(dim, dim) for dim in self.nesting_list])
 
         self.mlps = nn.ModuleList(
             [
@@ -197,7 +197,31 @@ class DINOLossMRL(nn.Module):
         teacher_out_softmax_centered_list: list[torch.Tensor],
         n_crops: int | tuple[int, int],
         teacher_global: bool,
-    ) -> float:
+    ) -> torch.Tensor:
+        total_loss = 0.0
+        if teacher_global is False:
+            for student_outputs, teacher_outputs in zip(student_output_list, teacher_out_softmax_centered_list):
+                s = torch.stack(student_outputs.chunk(n_crops[0]), 0)  # type: ignore[index]
+                t = teacher_outputs.view(n_crops[1], -1, teacher_outputs.shape[-1])  # type: ignore[index]
+                lsm = F.log_softmax(s / self.student_temp, dim=-1)
+                total_loss -= torch.einsum("tbk,sbk->tsb", t, lsm).mean(-1).sum()
+
+        else:
+            for student_outputs, teacher_outputs in zip(student_output_list, teacher_out_softmax_centered_list):
+                teacher_outputs = teacher_outputs.view(n_crops, -1, teacher_outputs.shape[-1])
+                lsm = F.log_softmax(student_outputs / self.student_temp, dim=-1)
+                loss = torch.sum(teacher_outputs.flatten(0, 1) * lsm, dim=-1)
+                total_loss -= loss.mean()
+
+        return total_loss
+
+    def forward_reference(
+        self,
+        student_output_list: list[torch.Tensor],
+        teacher_out_softmax_centered_list: list[torch.Tensor],
+        n_crops: int | tuple[int, int],
+        teacher_global: bool,
+    ) -> torch.Tensor:
         total_loss = 0.0
         if teacher_global is False:
             for student_outputs, teacher_outputs in zip(student_output_list, teacher_out_softmax_centered_list):

birder/net/ssl/i_jepa.py

@@ -69,11 +69,11 @@ class MultiBlockMasking:
     ) -> tuple[int, int]:
         _rand = torch.rand(1).item()
 
-        (min_s, max_s) = scale
+        min_s, max_s = scale
         mask_scale = min_s + _rand * (max_s - min_s)
         max_keep = int(self.height * self.width * mask_scale)
 
-        (min_ar, max_ar) = aspect_ratio_scale
+        min_ar, max_ar = aspect_ratio_scale
         aspect_ratio = min_ar + _rand * (max_ar - min_ar)
 
         # Compute block height and width (given scale and aspect-ratio)
@@ -154,7 +154,7 @@ class MultiBlockMasking:
             masks_p = []
             masks_c = []
             for _ in range(self.n_pred):
-                (mask, mask_c) = self._sample_block_mask(p_size)
+                mask, mask_c = self._sample_block_mask(p_size)
                 masks_p.append(mask)
                 masks_c.append(mask_c)
                 min_keep_pred = min(min_keep_pred, len(mask))
@@ -167,7 +167,7 @@ class MultiBlockMasking:
 
             masks_e = []
             for _ in range(self.n_enc):
-                (mask, _) = self._sample_block_mask(e_size, acceptable_regions=acceptable_regions)
+                mask, _ = self._sample_block_mask(e_size, acceptable_regions=acceptable_regions)
                 masks_e.append(mask)
                 min_keep_enc = min(min_keep_enc, len(mask))
 

birder/net/ssl/mmcr.py

@@ -125,7 +125,7 @@ class MMCR(SSLBaseNet):
         self.momentum_encoder.load_state_dict(self.encoder.state_dict())
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        (C, H, W) = x.shape[-3:]  # B, num_views, C, H, W
+        C, H, W = x.shape[-3:]  # B, num_views, C, H, W
         x = x.reshape(-1, C, H, W)
         z = self.encoder(x)
 

birder/net/swiftformer.py

@@ -111,7 +111,7 @@ class EfficientAdditiveAttention(nn.Module):
         self.final = nn.Linear(token_dim * num_heads, token_dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, _, H, W) = x.size()
+        B, _, H, W = x.size()
         x = x.flatten(2).permute(0, 2, 1)
 
         query = F.normalize(self.to_query(x), dim=-1)

birder/net/swin_transformer_v1.py

@@ -30,7 +30,7 @@ from birder.net.base import DetectorBackbone
 
 
 def patch_merging_pad(x: torch.Tensor) -> torch.Tensor:
-    (H, W, _) = x.shape[-3:]
+    H, W, _ = x.shape[-3:]
     x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
     x0 = x[..., 0::2, 0::2, :]  # ... H/2 W/2 C
     x1 = x[..., 1::2, 0::2, :]  # ... H/2 W/2 C
@@ -73,13 +73,13 @@ def shifted_window_attention(
     proj_bias: Optional[torch.Tensor] = None,
     logit_scale: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    (B, H, W, C) = x.size()
+    B, H, W, C = x.size()
 
     # Pad feature maps to multiples of window size
     pad_b = (window_size[0] - H % window_size[0]) % window_size[0]
     pad_r = (window_size[1] - W % window_size[1]) % window_size[1]
     x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
-    (_, pad_h, pad_w, _) = x.size()
+    _, pad_h, pad_w, _ = x.size()
 
     # If window size is larger than feature size, there is no need to shift window
     shift_size_w = shift_size[0]
@@ -309,7 +309,6 @@ class Swin_Transformer_v1(DetectorBackbone):
                 kernel_size=(patch_size, patch_size),
                 stride=(patch_size, patch_size),
                 padding=(0, 0),
-                bias=True,
             ),
             Permute([0, 2, 3, 1]),
             nn.LayerNorm(embed_dim, eps=1e-5),
@@ -434,7 +433,7 @@ class Swin_Transformer_v1(DetectorBackbone):
                     num_attn_heads = rel_pos_bias.size(1)
 
                     def _calc(src: int, dst: int) -> list[float]:
-                        (left, right) = 1.01, 1.5
+                        left, right = 1.01, 1.5
                         while right - left > 1e-6:
                             q = (left + right) / 2.0
                             gp = (1.0 - q ** (src // 2)) / (1.0 - q)  # Geometric progression

birder/net/swin_transformer_v2.py

@@ -76,7 +76,9 @@ class ShiftedWindowAttention(nn.Module):
 
         # MLP to generate continuous relative position bias
         self.cpb_mlp = nn.Sequential(
-            nn.Linear(2, 512, bias=True), nn.ReLU(inplace=True), nn.Linear(512, num_heads, bias=False)
+            nn.Linear(2, 512),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, num_heads, bias=False),
         )
         if qkv_bias is True:
             length = self.qkv.bias.numel() // 3
@@ -224,12 +226,7 @@ class Swin_Transformer_v2(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentio
 
         self.stem = nn.Sequential(
             nn.Conv2d(
-                self.input_channels,
-                embed_dim,
-                kernel_size=(patch_size, patch_size),
-                stride=patch_size,
-                padding=(0, 0),
-                bias=True,
+                self.input_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, padding=(0, 0)
             ),
             Permute([0, 2, 3, 1]),
             nn.LayerNorm(embed_dim, eps=1e-5),

birder/net/tiny_vit.py

@@ -201,12 +201,12 @@ class Attention(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         attn_bias = self.attention_biases[:, self.attention_bias_idxs]
-        (B, N, _) = x.shape
+        B, N, _ = x.shape
 
         # Normalization
         x = self.norm(x)
         qkv = self.qkv(x)
-        (q, k, v) = qkv.view(B, N, self.num_heads, -1).split([self.key_dim, self.key_dim, self.val_dim], dim=3)
+        q, k, v = qkv.view(B, N, self.num_heads, -1).split([self.key_dim, self.key_dim, self.val_dim], dim=3)
 
         q = q.permute(0, 2, 1, 3)
         k = k.permute(0, 2, 1, 3)
@@ -252,7 +252,7 @@ class TinyVitBlock(nn.Module):
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, H, W, C) = x.shape
+        B, H, W, C = x.shape
         L = H * W
 
         shortcut = x

birder/net/transnext.py

@@ -32,8 +32,8 @@ def get_relative_position_cpb(
     axis_kh = F.adaptive_avg_pool1d(axis_qh.unsqueeze(0), key_size[0]).squeeze(0)  # pylint: disable=not-callable
     axis_qw = torch.arange(query_size[1], dtype=torch.float32, device=device)
     axis_kw = F.adaptive_avg_pool1d(axis_qw.unsqueeze(0), key_size[1]).squeeze(0)  # pylint: disable=not-callable
-    (axis_kh, axis_kw) = torch.meshgrid(axis_kh, axis_kw, indexing="ij")
-    (axis_qh, axis_qw) = torch.meshgrid(axis_qh, axis_qw, indexing="ij")
+    axis_kh, axis_kw = torch.meshgrid(axis_kh, axis_kw, indexing="ij")
+    axis_qh, axis_qw = torch.meshgrid(axis_qh, axis_qw, indexing="ij")
 
     axis_kh = torch.reshape(axis_kh, [-1])
     axis_kw = torch.reshape(axis_kw, [-1])
@@ -44,7 +44,7 @@ def get_relative_position_cpb(
     relative_w = (axis_qw[:, None] - axis_kw[None, :]) / (pretrain_size[1] - 1) * 8
     relative_hw = torch.stack([relative_h, relative_w], dim=-1).view(-1, 2)
 
-    (relative_coords_table, idx_map) = torch.unique(relative_hw, return_inverse=True, dim=0)
+    relative_coords_table, idx_map = torch.unique(relative_hw, return_inverse=True, dim=0)
 
     relative_coords_table = (
         torch.sign(relative_coords_table)
@@ -86,9 +86,9 @@ class ConvolutionalGLU(nn.Module):
         self.drop = nn.Dropout(drop)
 
     def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
-        (x, v) = self.fc1(x).chunk(2, dim=-1)
+        x, v = self.fc1(x).chunk(2, dim=-1)
 
-        (B, _, C) = x.size()
+        B, _, C = x.size()
         x = x.transpose(1, 2).view(B, C, H, W).contiguous()
         x = self.dwconv(x)
         x = x.flatten(2).transpose(1, 2)
@@ -143,9 +143,9 @@ class Attention(nn.Module):
     def forward(
         self, x: torch.Tensor, _h: int, _w: int, relative_pos_index: torch.Tensor, relative_coords_table: torch.Tensor
     ) -> torch.Tensor:
-        (B, N, C) = x.size()
+        B, N, C = x.size()
         qkv = self.qkv(x).reshape(B, -1, 3 * self.num_heads, self.head_dim).permute(0, 2, 1, 3)
-        (q, k, v) = qkv.chunk(3, dim=1)
+        q, k, v = qkv.chunk(3, dim=1)
 
         # Use MLP to generate continuous relative positional bias
         rel_bias = (
@@ -217,9 +217,9 @@ class AggregatedAttention(nn.Module):
         self.act = nn.GELU()
 
         # MLP to generate continuous relative position bias
-        self.cpb_fc1 = nn.Linear(2, 512, bias=True)
+        self.cpb_fc1 = nn.Linear(2, 512)
         self.cpb_act = nn.ReLU(inplace=True)
-        self.cpb_fc2 = nn.Linear(512, num_heads, bias=True)
+        self.cpb_fc2 = nn.Linear(512, num_heads)
 
         # relative bias for local features
         self.relative_pos_bias_local = nn.Parameter(
@@ -227,7 +227,7 @@ class AggregatedAttention(nn.Module):
         )
 
         # Generate padding_mask and sequence length scale
-        (local_seq_length, padding_mask) = get_seqlen_and_mask(input_resolution, self.window_size)
+        local_seq_length, padding_mask = get_seqlen_and_mask(input_resolution, self.window_size)
         self.seq_length_scale = nn.Buffer(torch.log(local_seq_length + self.pool_len), persistent=False)
         self.padding_mask = nn.Buffer(padding_mask, persistent=False)
 
@@ -240,7 +240,7 @@ class AggregatedAttention(nn.Module):
     def forward(
         self, x: torch.Tensor, H: int, W: int, relative_pos_index: torch.Tensor, relative_coords_table: torch.Tensor
     ) -> torch.Tensor:
-        (B, N, C) = x.size()
+        B, N, C = x.size()
 
         # Generate queries, normalize them with L2, add query embedding,
         # and then magnify with sequence length scale and temperature.
@@ -252,7 +252,7 @@ class AggregatedAttention(nn.Module):
             * self.seq_length_scale
         )
 
-        (attn_local, v_local) = self.swa_qk_rpb(
+        attn_local, v_local = self.swa_qk_rpb(
             self.kv(x),
             q_norm_scaled.contiguous(),
             self.relative_pos_bias_local,
@@ -272,7 +272,7 @@ class AggregatedAttention(nn.Module):
 
         # Generate pooled keys and values
         kv_pool = self.kv(x_).reshape(B, self.pool_len, 2 * self.num_heads, self.head_dim).permute(0, 2, 1, 3)
-        (k_pool, v_pool) = kv_pool.chunk(2, dim=1)
+        k_pool, v_pool = kv_pool.chunk(2, dim=1)
 
         # Use MLP to generate continuous relative positional bias for pooled features.
         pool_bias = (
@@ -288,7 +288,7 @@ class AggregatedAttention(nn.Module):
         attn = self.attn_drop(attn)
 
         # Split the attention weights and separately aggregate the values of local & pooled features
-        (attn_local, attn_pool) = torch.split(attn, [self.local_len, self.pool_len], dim=-1)
+        attn_local, attn_pool = torch.split(attn, [self.local_len, self.pool_len], dim=-1)
 
         x_local = self.swa_av(
             q_norm, attn_local, v_local.contiguous(), self.learnable_tokens, self.learnable_bias, self.window_size, H, W
@@ -367,7 +367,7 @@ class OverlapPatchEmbed(nn.Module):
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int]:
         x = self.proj(x)
-        (_, _, H, W) = x.size()
+        _, _, H, W = x.size()
         x = x.flatten(2).transpose(1, 2)
         x = self.norm(x)
 
@@ -396,7 +396,7 @@ class TransNeXtStage(nn.Module):
 
         # Generate relative positional coordinate table and index for each stage
         # to compute continuous relative positional bias
-        (relative_pos_index, relative_coords_table) = get_relative_position_cpb(
+        relative_pos_index, relative_coords_table = get_relative_position_cpb(
             query_size=input_resolution, key_size=(input_resolution[0] // sr_ratio, input_resolution[1] // sr_ratio)
         )
         self.relative_pos_index = nn.Buffer(relative_pos_index, persistent=False)
@@ -430,7 +430,7 @@ class TransNeXtStage(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         B = x.size(0)
-        (x, H, W) = self.patch_embed(x)
+        x, H, W = self.patch_embed(x)
         for blk in self.blocks:
             x = blk(x, H, W, self.relative_pos_index, self.relative_coords_table)
 
@@ -553,7 +553,7 @@ class TransNeXt(DetectorBackbone):
                 sr_ratio = self.sr_ratio[i]
                 with torch.no_grad():
                     device = next(m.parameters()).device
-                    (relative_pos_index, relative_coords_table) = get_relative_position_cpb(
+                    relative_pos_index, relative_coords_table = get_relative_position_cpb(
                         query_size=input_resolution,
                         key_size=(input_resolution[0] // sr_ratio, input_resolution[1] // sr_ratio),
                         device=device,
@@ -574,7 +574,7 @@ class TransNeXt(DetectorBackbone):
                             blk.pool_len = pool_h * pool_w
                             blk.pool = nn.AdaptiveAvgPool2d((pool_h, pool_w))
 
-                            (local_seq_length, padding_mask) = get_seqlen_and_mask(
+                            local_seq_length, padding_mask = get_seqlen_and_mask(
                                 input_resolution, blk.window_size, device=device
                             )
                             blk.seq_length_scale = nn.Buffer(

birder/net/uniformer.py

@@ -71,9 +71,9 @@ class Attention(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, k, v) = qkv.unbind(0)
+        q, k, v = qkv.unbind(0)
 
         x = F.scaled_dot_product_attention(  # pylint: disable=not-callable
             q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0, scale=self.scale
@@ -137,7 +137,7 @@ class AttentionBlock(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = x + self.pos_embed(x)
-        (B, N, H, W) = x.shape
+        B, N, H, W = x.shape
         x = x.flatten(2).transpose(1, 2)
         x = x + self.drop_path(self.layer_scale_1(self.attn(self.norm1(x))))
         x = x + self.drop_path(self.layer_scale_2(self.mlp(self.norm2(x))))
@@ -155,7 +155,7 @@ class PatchEmbed(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.proj(x)
-        (B, _, H, W) = x.size()  # B, C, H, W
+        B, _, H, W = x.size()  # B, C, H, W
         x = x.flatten(2).transpose(1, 2)
         x = self.norm(x)
         x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()

birder/net/vgg.py

@@ -40,16 +40,7 @@ class Vgg(DetectorBackbone):
                 else:
                     in_channels = filters[i]
 
-                layers.append(
-                    nn.Conv2d(
-                        in_channels,
-                        filters[i],
-                        kernel_size=(3, 3),
-                        stride=(1, 1),
-                        padding=(1, 1),
-                        bias=True,
-                    )
-                )
+                layers.append(nn.Conv2d(in_channels, filters[i], kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))
                 layers.append(nn.ReLU(inplace=True))
 
             layers.append(nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)))