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

birder/net/vit.py

@@ -40,6 +40,7 @@ from birder.net.base import MaskedTokenRetentionMixin
 from birder.net.base import PreTrainEncoder
 from birder.net.base import TokenOmissionResultType
 from birder.net.base import TokenRetentionResultType
+from birder.net.base import normalize_out_indices
 
 
 def adjust_position_embedding(
@@ -73,12 +74,10 @@ def adjust_position_embedding(
 class PatchEmbed(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
-        The entire forward is equivalent to x.flatten(2).transpose(1, 2)
+        This is equivalent (in output) to: x.flatten(2).transpose(1, 2)
         """
 
-        (n, hidden_dim, h, w) = x.size()
-
-        # (n, hidden_dim, h, w) -> (n, hidden_dim, (h * w))
+        n, hidden_dim, h, w = x.size()
         x = x.reshape(n, hidden_dim, h * w)
 
         # (n, hidden_dim, (h * w)) -> (n, (h * w), hidden_dim)
@@ -155,9 +154,9 @@ class Attention(nn.Module):
             - attn_weights: If need_weights is True attention weights, otherwise, None.
         """
 
-        (B, N, C) = x.size()
+        B, N, C = x.size()
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
-        (q, k, v) = qkv.unbind(0)
+        q, k, v = qkv.unbind(0)
         q = self.q_norm(q)
         k = self.k_norm(k)
 
@@ -245,7 +244,7 @@ class EncoderBlock(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # torch._assert(x.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {x.size()}")
-        (attn_out, _) = self.attn(
+        attn_out, _ = self.attn(
             self.norm1(x),
             need_weights=self.need_attn,
             average_attn_weights=False,
@@ -317,13 +316,15 @@ class Encoder(nn.Module):
         x = self.pre_block(x)
         return self.block(x)
 
-    def forward_features(self, x: torch.Tensor) -> list[torch.Tensor]:
+    def forward_features(self, x: torch.Tensor, out_indices: Optional[list[int]] = None) -> list[torch.Tensor]:
         x = self.pre_block(x)
 
+        out_indices_set = set(out_indices) if out_indices is not None else None
         xs = []
-        for blk in self.block:
+        for idx, blk in enumerate(self.block):
             x = blk(x)
-            xs.append(x)
+            if out_indices_set is None or idx in out_indices_set:
+                xs.append(x)
 
         return xs
 
@@ -340,7 +341,7 @@ class Encoder(nn.Module):
 class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTokenRetentionMixin):
     block_group_regex = r"encoder\.block\.(\d+)"
 
-    # pylint: disable=too-many-locals,too-many-branches
+    # pylint: disable=too-many-locals,too-many-branches,too-many-statements
     def __init__(
         self,
         input_channels: int,
@@ -375,6 +376,7 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
         norm_layer_eps: float = self.config.get("norm_layer_eps", 1e-6)
         mlp_layer_type: str = self.config.get("mlp_layer_type", "FFN")
         act_layer_type: Optional[str] = self.config.get("act_layer_type", None)  # Default according to mlp type
+        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
         drop_path_rate: float = self.config["drop_path_rate"]
 
         if norm_layer_type == "LayerNorm":
@@ -405,6 +407,7 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
         self.hidden_dim = hidden_dim
         self.num_reg_tokens = num_reg_tokens
         self.attn_pool_special_tokens = attn_pool_special_tokens
+        self.out_indices = normalize_out_indices(out_indices, num_layers)
         dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]  # Stochastic depth decay rule
 
         self.conv_proj = nn.Conv2d(
@@ -472,8 +475,9 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
 
             self.attn_pool = MultiHeadAttentionPool(hidden_dim, attn_pool_num_heads, mlp_dim, qkv_bias=True)
 
-        self.return_stages = ["neck"]  # Actually meaningless, just for completeness
-        self.return_channels = [hidden_dim]
+        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 = [hidden_dim] * num_return_stages
         self.embedding_size = hidden_dim
         self.classifier = self.create_classifier()
 
@@ -537,8 +541,12 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)
 
+    def transform_to_backbone(self) -> None:
+        super().transform_to_backbone()
+        self.norm = nn.Identity()
+
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
         x = self.conv_proj(x)
         x = self.patch_embed(x)
 
@@ -558,15 +566,20 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
         if self.pos_embed_special_tokens is True:
             x = x + self._get_pos_embed(H, W)
 
-        x = self.encoder(x)
-        x = self.norm(x)
+        if self.out_indices is None:
+            xs = [self.encoder(x)]
+        else:
+            xs = self.encoder.forward_features(x, out_indices=self.out_indices)
 
-        x = x[:, self.num_special_tokens :]
-        x = x.permute(0, 2, 1)
-        (B, C, _) = x.size()
-        x = x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+        out: dict[str, torch.Tensor] = {}
+        for stage_name, stage_x in zip(self.return_stages, xs):
+            stage_x = stage_x[:, self.num_special_tokens :]
+            stage_x = stage_x.permute(0, 2, 1)
+            B, C, _ = stage_x.size()
+            stage_x = stage_x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+            out[stage_name] = stage_x
 
-        return {self.return_stages[0]: x}
+        return out
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
@@ -589,7 +602,7 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
         return_all_features: bool = False,
         return_keys: Literal["all", "tokens", "embedding"] = "tokens",
     ) -> TokenOmissionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)
@@ -663,7 +676,7 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
         mask_token: Optional[torch.Tensor] = None,
         return_keys: Literal["all", "features", "embedding"] = "features",
     ) -> TokenRetentionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         x = self.conv_proj(x)
         x = mask_tensor(x, mask, mask_token=mask_token, patch_factor=self.max_stride // self.stem_stride)
@@ -694,7 +707,7 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
         if return_keys in ("all", "features"):
             features = x[:, self.num_special_tokens :]
             features = features.permute(0, 2, 1)
-            (B, C, _) = features.size()
+            B, C, _ = features.size()
             features = features.reshape(B, C, H // self.patch_size, W // self.patch_size)
             result["features"] = features
 
@@ -714,7 +727,7 @@ class ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedTok
         return result
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)

birder/net/vit_parallel.py

@@ -31,6 +31,7 @@ from birder.net.base import MaskedTokenRetentionMixin
 from birder.net.base import PreTrainEncoder
 from birder.net.base import TokenOmissionResultType
 from birder.net.base import TokenRetentionResultType
+from birder.net.base import normalize_out_indices
 from birder.net.vit import PatchEmbed
 from birder.net.vit import adjust_position_embedding
 
@@ -51,9 +52,9 @@ class Attention(nn.Module):
         self.proj_drop = nn.Dropout(proj_drop)
 
     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, self.head_dim).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, is_causal=self.is_causal, scale=self.scale
@@ -172,11 +173,13 @@ class Encoder(nn.Module):
 
         return x
 
-    def forward_features(self, x: torch.Tensor) -> list[torch.Tensor]:
+    def forward_features(self, x: torch.Tensor, out_indices: Optional[list[int]] = None) -> list[torch.Tensor]:
         xs = []
-        for blk in self.block:
+        out_indices_set = set(out_indices) if out_indices is not None else None
+        for idx, blk in enumerate(self.block):
             x = blk(x)
-            xs.append(x)
+            if out_indices_set is None or idx in out_indices_set:
+                xs.append(x)
 
         return xs
 
@@ -213,6 +216,7 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         num_reg_tokens: int = self.config.get("num_reg_tokens", 0)
         class_token: bool = self.config.get("class_token", True)
         norm_layer_type: str = self.config.get("norm_layer_type", "LayerNorm")
+        out_indices: Optional[list[int]] = self.config.get("out_indices", None)
         drop_path_rate: float = self.config["drop_path_rate"]
 
         if norm_layer_type == "LayerNorm":
@@ -230,6 +234,7 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         self.hidden_dim = hidden_dim
         self.layer_scale_init_value = layer_scale_init_value
         self.num_reg_tokens = num_reg_tokens
+        self.out_indices = normalize_out_indices(out_indices, num_layers)
         dpr = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]  # Stochastic depth decay rule
 
         self.conv_proj = nn.Conv2d(
@@ -238,7 +243,6 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             kernel_size=(patch_size, patch_size),
             stride=(patch_size, patch_size),
             padding=(0, 0),
-            bias=True,
         )
         self.patch_embed = PatchEmbed()
 
@@ -278,8 +282,9 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         )
         self.norm = norm_layer(hidden_dim, eps=1e-6)
 
-        self.return_stages = ["neck"]  # Actually meaningless, but for completeness
-        self.return_channels = [hidden_dim]
+        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 = [hidden_dim] * num_return_stages
         self.embedding_size = hidden_dim
         self.classifier = self.create_classifier()
 
@@ -338,8 +343,12 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)
 
+    def transform_to_backbone(self) -> None:
+        super().transform_to_backbone()
+        self.norm = nn.Identity()
+
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
         x = self.conv_proj(x)
         x = self.patch_embed(x)
 
@@ -354,15 +363,21 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             x = torch.concat([batch_reg_tokens, x], dim=1)
 
         x = x + self._get_pos_embed(H, W)
-        x = self.encoder(x)
-        x = self.norm(x)
 
-        x = x[:, self.num_special_tokens :]
-        x = x.permute(0, 2, 1)
-        (B, C, _) = x.size()
-        x = x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+        if self.out_indices is None:
+            xs = [self.encoder(x)]
+        else:
+            xs = self.encoder.forward_features(x, out_indices=self.out_indices)
+
+        out: dict[str, torch.Tensor] = {}
+        for stage_name, stage_x in zip(self.return_stages, xs):
+            stage_x = stage_x[:, self.num_special_tokens :]
+            stage_x = stage_x.permute(0, 2, 1)
+            B, C, _ = stage_x.size()
+            stage_x = stage_x.reshape(B, C, H // self.patch_size, W // self.patch_size)
+            out[stage_name] = stage_x
 
-        return {self.return_stages[0]: x}
+        return out
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
@@ -384,7 +399,7 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         return_all_features: bool = False,
         return_keys: Literal["all", "tokens", "embedding"] = "tokens",
     ) -> TokenOmissionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)
@@ -441,7 +456,7 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         mask_token: Optional[torch.Tensor] = None,
         return_keys: Literal["all", "features", "embedding"] = "features",
     ) -> TokenRetentionResultType:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         x = self.conv_proj(x)
         x = mask_tensor(x, mask, mask_token=mask_token, patch_factor=self.max_stride // self.stem_stride)
@@ -467,7 +482,7 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         if return_keys in ("all", "features"):
             features = x[:, self.num_special_tokens :]
             features = features.permute(0, 2, 1)
-            (B, C, _) = features.size()
+            B, C, _ = features.size()
             features = features.reshape(B, C, H // self.patch_size, W // self.patch_size)
             result["features"] = features
 
@@ -481,7 +496,7 @@ class ViT_Parallel(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         return result
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
-        (H, W) = x.shape[-2:]
+        H, W = x.shape[-2:]
 
         # Reshape and permute the input tensor
         x = self.conv_proj(x)

birder/net/vit_sam.py

@@ -35,7 +35,7 @@ from birder.net.vit import EncoderBlock as MAEDecoderBlock
 
 # pylint: disable=invalid-name
 def window_partition(x: torch.Tensor, window_size: int) -> tuple[torch.Tensor, tuple[int, int]]:
-    (B, H, W, C) = x.shape
+    B, H, W, C = x.shape
 
     pad_h = (window_size - H % window_size) % window_size
     pad_w = (window_size - W % window_size) % window_size
@@ -55,8 +55,8 @@ def window_partition(x: torch.Tensor, window_size: int) -> tuple[torch.Tensor, t
 def window_unpartition(
     windows: torch.Tensor, window_size: int, pad_hw: tuple[int, int], hw: tuple[int, int]
 ) -> torch.Tensor:
-    (Hp, Wp) = pad_hw
-    (H, W) = hw
+    Hp, Wp = pad_hw
+    H, W = hw
     B = windows.shape[0] // (Hp * Wp // window_size // window_size)
     x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
     x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)
@@ -91,12 +91,12 @@ def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor
 def get_decomposed_rel_pos_bias(
     q: torch.Tensor, rel_pos_h: torch.Tensor, rel_pos_w: torch.Tensor, q_size: tuple[int, int], k_size: tuple[int, int]
 ) -> torch.Tensor:
-    (q_h, q_w) = q_size
-    (k_h, k_w) = k_size
+    q_h, q_w = q_size
+    k_h, k_w = k_size
     Rh = get_rel_pos(q_h, k_h, rel_pos_h)
     Rw = get_rel_pos(q_w, k_w, rel_pos_w)
 
-    (B, _, dim) = q.shape
+    B, _, dim = q.shape
     r_q = q.reshape(B, q_h, q_w, dim)
     rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)
     rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)
@@ -139,9 +139,9 @@ class Attention(nn.Module):
             self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, H, W, _) = x.shape
+        B, H, W, _ = x.shape
         qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
-        (q, k, v) = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
+        q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
 
         if self.use_rel_pos is True:
             attn_bias = get_decomposed_rel_pos_bias(q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))
@@ -216,13 +216,13 @@ class EncoderBlock(nn.Module):
             self.layer_scale_2 = nn.Identity()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (_, H, W, _) = x.shape
+        _, H, W, _ = x.shape
         shortcut = x
 
         x = self.norm1(x)
         pad_hw = (0, 0)
         if self.window_size > 0:
-            (x, pad_hw) = window_partition(x, self.window_size)
+            x, pad_hw = window_partition(x, self.window_size)
 
         x = self.attn(x)
         if self.window_size > 0:

birder/net/vovnet_v2.py

@@ -27,7 +27,7 @@ class EffectiveSE(nn.Module):
 
     def __init__(self, channels: int) -> None:
         super().__init__()
-        self.fc = nn.Conv2d(channels, channels, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=True)
+        self.fc = nn.Conv2d(channels, channels, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x_se = x.mean(dim=(2, 3), keepdim=True)

birder/net/xcit.py

@@ -30,6 +30,7 @@ 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
 from birder.net.cait import ClassAttention
 
 
@@ -212,7 +213,7 @@ class LPI(nn.Module):
         )
 
     def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
-        (B, N, C) = x.shape
+        B, N, C = x.shape
         x = x.permute(0, 2, 1).reshape(B, C, H, W)
         x = self.conv_bn_act(x)
         x = self.conv(x)
@@ -236,10 +237,10 @@ class XCA(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)
         qkv = qkv.permute(2, 0, 3, 1, 4)
-        (q, k, v) = qkv.unbind(0)
+        q, k, v = qkv.unbind(0)
 
         q = F.normalize(q, dim=-1) * self.temperature
         k = F.normalize(k, dim=-1)
@@ -311,6 +312,7 @@ class XCiT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
         else:
             raise ValueError(f"depth={depth} is not supported")
 
+        out_indices = normalize_out_indices(out_indices, depth)
         self.patch_embed = ConvPatchEmbed(patch_size, self.input_channels, dim=embed_dim)
         self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
         dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # Stochastic depth decay rule
@@ -381,7 +383,7 @@ class XCiT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
     def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
         B = x.size(0)
 
-        (x, H, W) = self.patch_embed(x)
+        x, H, W = self.patch_embed(x)
 
         pos_encoding = self.pos_embed(B, H, W).reshape(B, -1, x.size(1)).permute(0, 2, 1)
         x = x + pos_encoding
@@ -414,7 +416,7 @@ class XCiT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
     ) -> TokenRetentionResultType:
         B = x.size(0)
 
-        (x, H, W) = self.patch_embed(x)
+        x, H, W = self.patch_embed(x)
         x = mask_tensor(
             x.permute(0, 2, 1).reshape(B, -1, H, W),
             mask,
@@ -435,7 +437,7 @@ class XCiT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
         if return_keys in ("all", "features"):
             features = x[:, 1:]
             features = features.permute(0, 2, 1)
-            (B, C, _) = features.size()
+            B, C, _ = features.size()
             features = features.reshape(B, C, H, W)
             result["features"] = features
 
@@ -447,7 +449,7 @@ class XCiT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         B = x.size(0)
 
-        (x, H, W) = self.patch_embed(x)
+        x, H, W = self.patch_embed(x)
 
         pos_encoding = self.pos_embed(B, H, W).reshape(B, -1, x.size(1)).permute(0, 2, 1)
         x = x + pos_encoding

birder/ops/msda.py

@@ -91,8 +91,8 @@ def _ms_deform_attn_setup_context(  # type: ignore[no-untyped-def] # pylint: dis
 
 
 def _ms_deform_attn_backward(ctx, grad_output):  # type: ignore[no-untyped-def]
-    (value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights) = ctx.saved_tensors
-    (grad_value, grad_sampling_loc, grad_attn_weight) = ms_deform_attn_backward_op(
+    value, value_spatial_shapes, value_level_start_index, sampling_locations, attention_weights = ctx.saved_tensors
+    grad_value, grad_sampling_loc, grad_attn_weight = ms_deform_attn_backward_op(
         value,
         value_spatial_shapes,
         value_level_start_index,
@@ -160,8 +160,8 @@ def multi_scale_deformable_attention(
     attention_weights: torch.Tensor,
     im2col_step: int,  # pylint: disable=unused-argument
 ) -> torch.Tensor:
-    (batch_size, _, num_heads, hidden_dim) = value.size()
-    (_, num_queries, num_heads, num_levels, num_points, _) = sampling_locations.size()
+    batch_size, _, num_heads, hidden_dim = value.size()
+    _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.size()
     areas: list[int] = value_spatial_shapes.prod(dim=1).tolist()
     value_list = value.split(areas, dim=1)
     sampling_grids = 2 * sampling_locations - 1

birder/ops/swattention.py

@@ -38,7 +38,7 @@ def _swattention_qk_rpb_fake(  # pylint: disable=unused-argument
 def _swattention_qk_rpb_setup_context(  # type: ignore[no-untyped-def] # pylint: disable=unused-argument
     ctx, inputs, output
 ) -> None:
-    (query, key, _rpb, height, width, kernel_size) = inputs
+    query, key, _rpb, height, width, kernel_size = inputs
     ctx.save_for_backward(query, key)
     ctx.height = height
     ctx.width = width
@@ -46,8 +46,8 @@ def _swattention_qk_rpb_setup_context(  # type: ignore[no-untyped-def] # pylint:
 
 
 def _swattention_qk_rpb_backward(ctx, grad_output):  # type: ignore[no-untyped-def]
-    (query, key) = ctx.saved_tensors
-    (d_query, d_key, d_rpb) = swattention_qk_rpb_backward_op(
+    query, key = ctx.saved_tensors
+    d_query, d_key, d_rpb = swattention_qk_rpb_backward_op(
         grad_output.contiguous(), query, key, ctx.height, ctx.width, ctx.kernel_size
     )
     return (d_query, d_key, d_rpb, None, None, None)
@@ -107,8 +107,8 @@ def _swattention_av_setup_context(  # type: ignore[no-untyped-def] # pylint: dis
 
 
 def _swattention_av_backward(ctx, grad_output):  # type: ignore[no-untyped-def]
-    (attn_weight, value) = ctx.saved_tensors
-    (d_attn_weight, d_value) = swattention_av_backward_op(
+    attn_weight, value = ctx.saved_tensors
+    d_attn_weight, d_value = swattention_av_backward_op(
         grad_output.contiguous(), attn_weight, value, ctx.height, ctx.width, ctx.kernel_size
     )
     return (d_attn_weight, d_value, None, None, None)
@@ -184,10 +184,10 @@ class SWAttention_QK_RPB(nn.Module):
             )
 
         # Custom kernel
-        (B, N, _) = kv.size()
+        B, N, _ = kv.size()
 
         # Generate unfolded keys and values and l2-normalize them
-        (k_local, v_local) = kv.reshape(B, N, 2 * num_heads, head_dim).permute(0, 2, 1, 3).chunk(2, dim=1)
+        k_local, v_local = kv.reshape(B, N, 2 * num_heads, head_dim).permute(0, 2, 1, 3).chunk(2, dim=1)
 
         # Compute local similarity
         attn_local = swattention_qk_rpb_op(
@@ -254,14 +254,14 @@ def swattention_qk_rpb(
     H: int,
     W: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    (B, N, _) = kv.size()
+    B, N, _ = kv.size()
 
     # Generate unfolded keys and values and l2-normalize them
-    (k_local, v_local) = kv.chunk(2, dim=-1)
+    k_local, v_local = kv.chunk(2, dim=-1)
     k_local = F.normalize(k_local.reshape(B, N, num_heads, head_dim), dim=-1).reshape(B, N, -1)
     kv_local = torch.concat([k_local, v_local], dim=-1).permute(0, 2, 1).reshape(B, -1, H, W)
 
-    (k_local, v_local) = (
+    k_local, v_local = (
         F.unfold(kv_local, kernel_size=window_size, padding=window_size // 2, stride=1)
         .reshape(B, 2 * num_heads, head_dim, local_len, N)
         .permute(0, 1, 4, 2, 3)

birder/results/classification.py

@@ -30,7 +30,7 @@ def top_k_accuracy_score(y_true: npt.NDArray[Any], y_pred: npt.NDArray[np.float6
     if len(y_true.shape) == 2:
         y_true = np.argmax(y_true, axis=1)
 
-    (num_samples, _num_labels) = y_pred.shape
+    num_samples, _num_labels = y_pred.shape
     indices: list[int] = []
     arg_sorted = np.argpartition(y_pred, -top_k, axis=1)[:, -top_k:]
     for i in range(num_samples):
@@ -693,7 +693,7 @@ class SparseResults(Results):
             For sparse files, this value is ignored.
         """
 
-        (label_names, detected_sparse_k) = detect_file_format(path)
+        label_names, detected_sparse_k = detect_file_format(path)
 
         if detected_sparse_k is not None:
             schema_overrides = {
@@ -817,7 +817,7 @@ def load_results(path: str, lazy: bool = True) -> Results | SparseResults:
     <class 'birder.results.classification.SparseResults'>
     """
 
-    (_, sparse_k) = detect_file_format(path)
+    _, sparse_k = detect_file_format(path)
 
     # Load using appropriate class
     if sparse_k is not None:

birder/results/gui.py

@@ -212,7 +212,7 @@ class ConfusionMatrix:
             )
 
         offset = 0.5
-        (height, width) = cnf_matrix.shape
+        height, width = cnf_matrix.shape
         ax.hlines(
             y=np.arange(height + 1) - offset,
             xmin=-offset,
@@ -261,7 +261,7 @@ class ROC:
             roc_auc = {}
             for i in results.unique_labels:
                 binary_labels = results.labels == i
-                (fpr[i], tpr[i], _) = roc_curve(binary_labels, results.output[:, i])
+                fpr[i], tpr[i], _ = roc_curve(binary_labels, results.output[:, i])
                 if np.sum(binary_labels) == 0:
                     tpr[i] = np.zeros_like(fpr[i])
 
@@ -324,7 +324,7 @@ class PrecisionRecall:
             labels = label_binarize(results.labels, classes=range(len(results.label_names)))
 
             # A "micro-average" quantifying score on all classes jointly
-            (precision, recall, _) = precision_recall_curve(labels.ravel(), results.output.ravel())
+            precision, recall, _ = precision_recall_curve(labels.ravel(), results.output.ravel())
             average_precision = average_precision_score(labels.ravel(), results.output.ravel(), average="micro")
 
             line = ax.step(recall, precision, linestyle=":", where="post")
@@ -334,7 +334,7 @@ class PrecisionRecall:
             # Per selected class
             for cls in pr_classes:
                 i = results.label_names.index(cls)
-                (precision, recall, _) = precision_recall_curve(labels[:, i], results.output[:, i])
+                precision, recall, _ = precision_recall_curve(labels[:, i], results.output[:, i])
                 average_precision = average_precision_score(labels[:, i], results.output[:, i])
                 line = ax.plot(recall, precision, lw=2)
                 legend_lines.append(line[0])
@@ -372,8 +372,8 @@ class ProbabilityHistogram:
         cls_a_df = results_df.filter(pl.col("label_name") == cls_a)
         cls_b_df = results_df.filter(pl.col("label_name") == cls_b)
 
-        (cls_a_prob_a_counts, cls_a_prob_a_bins) = hist(cls_a_df[str(self.results.label_names.index(cls_a))])
-        (cls_a_prob_b_counts, cls_a_prob_b_bins) = hist(cls_b_df[str(self.results.label_names.index(cls_a))])
+        cls_a_prob_a_counts, cls_a_prob_a_bins = hist(cls_a_df[str(self.results.label_names.index(cls_a))])
+        cls_a_prob_b_counts, cls_a_prob_b_bins = hist(cls_b_df[str(self.results.label_names.index(cls_a))])
         plt.subplot(2, 1, 1)
         plt.stairs(
             cls_a_prob_a_counts,
@@ -391,8 +391,8 @@ class ProbabilityHistogram:
         )
         plt.legend(loc="upper center")
 
-        (cls_b_prob_a_counts, cls_b_prob_a_bins) = hist(cls_a_df[str(self.results.label_names.index(cls_b))])
-        (cls_b_prob_b_counts, cls_b_prob_b_bins) = hist(cls_b_df[str(self.results.label_names.index(cls_b))])
+        cls_b_prob_a_counts, cls_b_prob_a_bins = hist(cls_a_df[str(self.results.label_names.index(cls_b))])
+        cls_b_prob_b_counts, cls_b_prob_b_bins = hist(cls_b_df[str(self.results.label_names.index(cls_b))])
         plt.subplot(2, 1, 2)
         plt.stairs(
             cls_b_prob_b_counts,