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

birder/common/training_cli.py

@@ -56,7 +56,9 @@ def add_optimization_args(parser: argparse.ArgumentParser, default_batch_size: i
     )
 
 
-def add_lr_wd_args(parser: argparse.ArgumentParser, backbone_lr: bool = False, wd_end: bool = False) -> None:
+def add_lr_wd_args(
+    parser: argparse.ArgumentParser, backbone_lr: bool = False, wd_end: bool = False, backbone_layer_decay: bool = False
+) -> None:
     group = parser.add_argument_group("Learning rate and regularization parameters")
     group.add_argument("--lr", type=float, default=0.1, metavar="LR", help="base learning rate")
     group.add_argument("--bias-lr", type=float, metavar="LR", help="learning rate of biases")
@@ -92,6 +94,9 @@ def add_lr_wd_args(parser: argparse.ArgumentParser, backbone_lr: bool = False, w
         help="custom weight decay for specific layers by name (e.g., offset_conv=0.0)",
     )
     group.add_argument("--layer-decay", type=float, help="layer-wise learning rate decay (LLRD)")
+    if backbone_layer_decay is True:
+        group.add_argument("--backbone-layer-decay", type=float, help="backbone layer-wise learning rate decay (LLRD)")
+
     group.add_argument("--layer-decay-min-scale", type=float, help="minimum layer scale factor clamp value")
     group.add_argument(
         "--layer-decay-no-opt-scale", type=float, help="layer scale threshold below which parameters are frozen"

birder/common/training_utils.py

@@ -343,7 +343,7 @@ def count_layers(model: torch.nn.Module) -> int:
     return num_layers
 
 
-# pylint: disable=protected-access,too-many-locals,too-many-branches
+# pylint: disable=protected-access,too-many-locals,too-many-branches,too-many-statements
 def optimizer_parameter_groups(
     model: torch.nn.Module,
     weight_decay: float,
@@ -352,6 +352,7 @@ def optimizer_parameter_groups(
     custom_keys_weight_decay: Optional[list[tuple[str, float]]] = None,
     custom_layer_weight_decay: Optional[dict[str, float]] = None,
     layer_decay: Optional[float] = None,
+    backbone_layer_decay: Optional[float] = None,
     layer_decay_min_scale: Optional[float] = None,
     layer_decay_no_opt_scale: Optional[float] = None,
     bias_lr: Optional[float] = None,
@@ -388,6 +389,8 @@ def optimizer_parameter_groups(
         Applied to parameters whose names contain the specified keys.
     layer_decay
         Layer-wise learning rate decay factor.
+    backbone_layer_decay
+        Layer-wise learning rate decay factor for backbone parameters only.
     layer_decay_min_scale
         Minimum learning rate scale factor when using layer decay. Prevents layers from having too small learning rates.
     layer_decay_no_opt_scale
@@ -434,6 +437,27 @@ def optimizer_parameter_groups(
         if layer_decay is not None:
             logger.warning("Assigning lr scaling (layer decay) without a block group map")
 
+    backbone_group_map: dict[str, int] = {}
+    backbone_num_layers = 0
+    if backbone_layer_decay is not None:
+        backbone_module = getattr(model, "backbone", None)
+        if backbone_module is None:
+            logger.warning("Backbone layer decay requested but model has no backbone")
+            backbone_layer_decay = None
+        else:
+            backbone_block_group_regex = getattr(backbone_module, "block_group_regex", None)
+            if backbone_block_group_regex is not None:
+                names = [n for n, _ in backbone_module.named_parameters()]
+                groups = group_by_regex(names, backbone_block_group_regex)
+                backbone_group_map = {
+                    f"backbone.{item}": index for index, sublist in enumerate(groups) for item in sublist
+                }
+                backbone_num_layers = len(groups)
+            else:
+                backbone_group_map = {}
+                backbone_num_layers = count_layers(backbone_module)
+                logger.warning("Assigning lr scaling (backbone layer decay) without a block group map")
+
     # Build layer scale
     if layer_decay_min_scale is None:
         layer_decay_min_scale = 0.0
@@ -444,14 +468,28 @@ def optimizer_parameter_groups(
         layer_scales = [max(layer_decay_min_scale, layer_decay ** (layer_max - i)) for i in range(num_layers)]
         logger.info(f"Layer scaling ranges from {min(layer_scales)} to {max(layer_scales)} across {num_layers} layers")
 
+    backbone_layer_scales = []
+    if backbone_layer_decay is not None:
+        backbone_layer_max = backbone_num_layers - 1
+        backbone_layer_scales = [
+            max(layer_decay_min_scale, backbone_layer_decay ** (backbone_layer_max - i))
+            for i in range(backbone_num_layers)
+        ]
+        logger.info(
+            "Backbone layer scaling ranges from "
+            f"{min(backbone_layer_scales)} to {max(backbone_layer_scales)} across {backbone_num_layers} layers"
+        )
+
     # Set weight decay and layer decay
     idx = 0
+    backbone_idx = 0
     params = []
     module_stack_with_prefix = [(model, "")]
     visited_modules = []
     while len(module_stack_with_prefix) > 0:  # pylint: disable=too-many-nested-blocks
         skip_module = False
         module, prefix = module_stack_with_prefix.pop()
+        is_backbone_module = prefix == "backbone" or prefix.startswith("backbone.")
         if id(module) in visited_modules:
             skip_module = True
 
@@ -460,23 +498,35 @@ def optimizer_parameter_groups(
         for name, p in module.named_parameters(recurse=False):
             target_name = f"{prefix}.{name}" if prefix != "" else name
             idx = group_map.get(target_name, idx)
+            is_backbone_param = target_name.startswith("backbone.")
+            if backbone_layer_decay is not None and is_backbone_param is True:
+                backbone_idx = backbone_group_map.get(target_name, backbone_idx)
             if skip_module is True:
                 break
 
             parameters_found = True
             if p.requires_grad is False:
                 continue
-            if layer_decay is not None and layer_decay_no_opt_scale is not None:
-                if layer_scales[idx] < layer_decay_no_opt_scale:
-                    p.requires_grad_(False)
+            if layer_decay_no_opt_scale is not None:
+                if backbone_layer_decay is not None and is_backbone_param is True:
+                    if backbone_layer_scales and backbone_layer_scales[backbone_idx] < layer_decay_no_opt_scale:
+                        p.requires_grad_(False)
+                elif layer_decay is not None:
+                    if layer_scales[idx] < layer_decay_no_opt_scale:
+                        p.requires_grad_(False)
 
             is_custom_key = False
             if custom_keys_weight_decay is not None:
                 for key, custom_wd in custom_keys_weight_decay:
                     target_name_for_custom_key = f"{prefix}.{name}" if prefix != "" and "." in key else name
                     if key == target_name_for_custom_key:
-                        # Calculate lr_scale (from layer_decay or custom_layer_lr_scale)
-                        lr_scale = 1.0 if layer_decay is None else layer_scales[idx]
+                        # Calculate lr_scale (from layer_decay/backbone_layer_decay or custom_layer_lr_scale)
+                        if layer_decay is not None and (backbone_layer_decay is None or is_backbone_param is False):
+                            lr_scale = layer_scales[idx]
+                        elif backbone_layer_decay is not None and is_backbone_param is True:
+                            lr_scale = backbone_layer_scales[backbone_idx]
+                        else:
+                            lr_scale = 1.0
                         if custom_layer_lr_scale is not None:
                             for layer_name_key, custom_scale in custom_layer_lr_scale.items():
                                 if layer_name_key in target_name:
@@ -500,8 +550,8 @@ def optimizer_parameter_groups(
                         # Apply learning rate based on priority: bias_lr > backbone_lr > lr_scale
                         if bias_lr is not None and target_name.endswith(".bias") is True:
                             d["lr"] = bias_lr
-                        elif backbone_lr is not None and target_name.startswith("backbone.") is True:
-                            d["lr"] = backbone_lr
+                        elif backbone_lr is not None and is_backbone_param is True:
+                            d["lr"] = backbone_lr * lr_scale if backbone_layer_decay is not None else backbone_lr
                         elif lr_scale != 1.0:
                             d["lr"] = base_lr * lr_scale
 
@@ -522,8 +572,13 @@ def optimizer_parameter_groups(
                             wd = custom_wd_value
                             break
 
-                # Calculate lr_scale (from layer_decay or custom_layer_lr_scale)
-                lr_scale = 1.0 if layer_decay is None else layer_scales[idx]
+                # Calculate lr_scale (from layer_decay/backbone_layer_decay or custom_layer_lr_scale)
+                if layer_decay is not None and (backbone_layer_decay is None or is_backbone_param is False):
+                    lr_scale = layer_scales[idx]
+                elif backbone_layer_decay is not None and is_backbone_param is True:
+                    lr_scale = backbone_layer_scales[backbone_idx]
+                else:
+                    lr_scale = 1.0
                 if custom_layer_lr_scale is not None:
                     for layer_name_key, custom_scale in custom_layer_lr_scale.items():
                         if layer_name_key in target_name:
@@ -539,8 +594,8 @@ def optimizer_parameter_groups(
                 # Apply learning rate based on priority: bias_lr > backbone_lr > lr_scale
                 if bias_lr is not None and target_name.endswith(".bias") is True:
                     d["lr"] = bias_lr
-                elif backbone_lr is not None and target_name.startswith("backbone.") is True:
-                    d["lr"] = backbone_lr
+                elif backbone_lr is not None and is_backbone_param is True:
+                    d["lr"] = backbone_lr * lr_scale if backbone_layer_decay is not None else backbone_lr
                 elif lr_scale != 1.0:
                     d["lr"] = base_lr * lr_scale
 
@@ -548,6 +603,8 @@ def optimizer_parameter_groups(
 
         if parameters_found is True:
             idx += 1
+            if is_backbone_module is True:
+                backbone_idx += 1
 
         for child_name, child_module in reversed(list(module.named_children())):
             child_prefix = f"{prefix}.{child_name}" if prefix != "" else child_name

birder/net/_vit_configs.py

@@ -111,6 +111,11 @@ def register_vit_configs(vit: type[BaseNet]) -> None:
         vit,
         config={"patch_size": 16, **BASE, "layer_scale_init_value": 1e-5},
     )
+    registry.register_model_config(
+        "vit_b16_pn",
+        vit,
+        config={"patch_size": 16, **BASE, "pre_norm": True, "norm_layer_eps": 1e-5},
+    )
     registry.register_model_config(
         "vit_b16_qkn_ls",
         vit,

birder/net/cait.py

@@ -231,11 +231,11 @@ class CaiT(BaseNet):
             if isinstance(m, nn.Linear):
                 nn.init.trunc_normal_(m.weight, std=0.02)
                 if m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
+                    nn.init.zeros_(m.bias)
 
             elif isinstance(m, nn.LayerNorm):
-                nn.init.constant_(m.bias, 0)
-                nn.init.constant_(m.weight, 1.0)
+                nn.init.zeros_(m.bias)
+                nn.init.ones_(m.weight)
 
         nn.init.trunc_normal_(self.pos_embed, std=0.02)
         nn.init.trunc_normal_(self.cls_token, std=0.02)

birder/net/coat.py

@@ -474,11 +474,11 @@ class CoaT(DetectorBackbone):
             if isinstance(m, nn.Linear):
                 nn.init.trunc_normal_(m.weight, std=0.02)
                 if m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
+                    nn.init.zeros_(m.bias)
 
             elif isinstance(m, nn.LayerNorm):
-                nn.init.constant_(m.bias, 0)
-                nn.init.constant_(m.weight, 1.0)
+                nn.init.zeros_(m.bias)
+                nn.init.ones_(m.weight)
 
         nn.init.trunc_normal_(self.cls_token1, std=0.02)
         nn.init.trunc_normal_(self.cls_token2, std=0.02)

birder/net/deit.py

@@ -167,7 +167,7 @@ class DeiT(DetectorBackbone):
             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):
+        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
             stage_x = stage_x[:, self.num_special_tokens :]
             stage_x = stage_x.permute(0, 2, 1)
             B, C, _ = stage_x.size()

birder/net/deit3.py

@@ -185,7 +185,7 @@ class DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedT
             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):
+        for stage_name, stage_x in zip(self.return_stages, xs, strict=True):
             stage_x = stage_x[:, self.num_special_tokens :]
             stage_x = stage_x.permute(0, 2, 1)
             B, C, _ = stage_x.size()

birder/net/detection/__init__.py

@@ -3,6 +3,7 @@ from birder.net.detection.detr import DETR
 from birder.net.detection.efficientdet import EfficientDet
 from birder.net.detection.faster_rcnn import Faster_RCNN
 from birder.net.detection.fcos import FCOS
+from birder.net.detection.lw_detr import LW_DETR
 from birder.net.detection.plain_detr import Plain_DETR
 from birder.net.detection.retinanet import RetinaNet
 from birder.net.detection.rt_detr_v1 import RT_DETR_v1
@@ -21,6 +22,7 @@ __all__ = [
     "EfficientDet",
     "Faster_RCNN",
     "FCOS",
+    "LW_DETR",
     "Plain_DETR",
     "RetinaNet",
     "RT_DETR_v1",

birder/net/detection/deformable_detr.py

@@ -56,7 +56,7 @@ class HungarianMatcher(nn.Module):
     @torch.jit.unused  # type: ignore[untyped-decorator]
     def forward(
         self, class_logits: torch.Tensor, box_regression: torch.Tensor, targets: list[dict[str, torch.Tensor]]
-    ) -> list[torch.Tensor]:
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         with torch.no_grad():
             B, num_queries = class_logits.shape[:2]
 
@@ -135,7 +135,7 @@ class MultiScaleDeformableAttention(nn.Module):
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
-        nn.init.constant_(self.sampling_offsets.weight, 0.0)
+        nn.init.zeros_(self.sampling_offsets.weight)
         thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
         grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
         grid_init = (
@@ -149,12 +149,12 @@ class MultiScaleDeformableAttention(nn.Module):
         with torch.no_grad():
             self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
 
-        nn.init.constant_(self.attention_weights.weight, 0.0)
-        nn.init.constant_(self.attention_weights.bias, 0.0)
+        nn.init.zeros_(self.attention_weights.weight)
+        nn.init.zeros_(self.attention_weights.bias)
         nn.init.xavier_uniform_(self.value_proj.weight)
-        nn.init.constant_(self.value_proj.bias, 0.0)
+        nn.init.zeros_(self.value_proj.bias)
         nn.init.xavier_uniform_(self.output_proj.weight)
-        nn.init.constant_(self.output_proj.bias, 0.0)
+        nn.init.zeros_(self.output_proj.bias)
 
     def forward(
         self,
@@ -279,11 +279,10 @@ class DeformableTransformerDecoderLayer(nn.Module):
         self_attn_mask: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         # Self attention
-        q = tgt + query_pos
-        k = tgt + query_pos
+        q_k = tgt + query_pos
 
         tgt2, _ = self.self_attn(
-            q.transpose(0, 1), k.transpose(0, 1), tgt.transpose(0, 1), need_weights=False, attn_mask=self_attn_mask
+            q_k.transpose(0, 1), q_k.transpose(0, 1), tgt.transpose(0, 1), need_weights=False, attn_mask=self_attn_mask
         )
         tgt2 = tgt2.transpose(0, 1)
         tgt = tgt + self.dropout(tgt2)
@@ -587,7 +586,7 @@ class Deformable_DETR(DetectionBaseNet):
 
         self.query_embed = nn.Embedding(num_queries, hidden_dim * 2)
         self.pos_enc = PositionEmbeddingSine(hidden_dim // 2, normalize=True)
-        self.matcher = HungarianMatcher(cost_class=2, cost_bbox=5, cost_giou=2)
+        self.matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0)
 
         class_embed = nn.Linear(hidden_dim, self.num_classes)
         bbox_embed = MLP(hidden_dim, [hidden_dim, hidden_dim, 4], activation_layer=nn.ReLU)
@@ -641,7 +640,8 @@ class Deformable_DETR(DetectionBaseNet):
             for param in self.class_embed.parameters():
                 param.requires_grad_(True)
 
-    def _get_src_permutation_idx(self, indices: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    @staticmethod
+    def _get_src_permutation_idx(indices: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
         batch_idx = torch.concat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
         src_idx = torch.concat([src for (src, _) in indices])
         return (batch_idx, src_idx)
@@ -709,7 +709,7 @@ class Deformable_DETR(DetectionBaseNet):
         if training_utils.is_dist_available_and_initialized() is True:
             torch.distributed.all_reduce(num_boxes)
 
-        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1).item()
+        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1)
 
         loss_ce_list = []
         loss_bbox_list = []

birder/net/detection/detr.py

@@ -49,7 +49,7 @@ class HungarianMatcher(nn.Module):
     @torch.jit.unused  # type: ignore[untyped-decorator]
     def forward(
         self, class_logits: torch.Tensor, box_regression: torch.Tensor, targets: list[dict[str, torch.Tensor]]
-    ) -> list[torch.Tensor]:
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         with torch.no_grad():
             B, num_queries = class_logits.shape[:2]
 
@@ -148,10 +148,9 @@ class TransformerDecoderLayer(nn.Module):
         query_pos: torch.Tensor,
         memory_key_padding_mask: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        q = tgt + query_pos
-        k = tgt + query_pos
+        q_k = tgt + query_pos
 
-        tgt2, _ = self.self_attn(q, k, value=tgt, need_weights=False)
+        tgt2, _ = self.self_attn(q_k, q_k, value=tgt, need_weights=False)
         tgt = tgt + self.dropout1(tgt2)
         tgt = self.norm1(tgt)
         tgt2, _ = self.multihead_attn(
@@ -341,7 +340,7 @@ class DETR(DetectionBaseNet):
         )
         self.pos_enc = PositionEmbeddingSine(hidden_dim // 2, normalize=True)
 
-        self.matcher = HungarianMatcher(cost_class=1, cost_bbox=5, cost_giou=2)
+        self.matcher = HungarianMatcher(cost_class=1.0, cost_bbox=5.0, cost_giou=2.0)
         empty_weight = torch.ones(self.num_classes)
         empty_weight[0] = 0.1
         self.empty_weight = nn.Buffer(empty_weight)
@@ -365,7 +364,8 @@ class DETR(DetectionBaseNet):
             for param in self.class_embed.parameters():
                 param.requires_grad_(True)
 
-    def _get_src_permutation_idx(self, indices: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    @staticmethod
+    def _get_src_permutation_idx(indices: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
         batch_idx = torch.concat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
         src_idx = torch.concat([src for (src, _) in indices])
         return (batch_idx, src_idx)
@@ -422,7 +422,7 @@ class DETR(DetectionBaseNet):
         if training_utils.is_dist_available_and_initialized() is True:
             torch.distributed.all_reduce(num_boxes)
 
-        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1).item()
+        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1)
 
         loss_ce_list = []
         loss_bbox_list = []