birder 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

birder/common/lib.py

@@ -1,11 +1,7 @@
 import os
-import random
 from typing import Any
 from typing import Optional
 
-import numpy as np
-import torch
-
 from birder.conf import settings
 from birder.data.transforms.classification import RGBType
 from birder.model_registry import registry
@@ -19,11 +15,8 @@ from birder.net.ssl.base import SSLBaseNet
 from birder.version import __version__
 
 
-def set_random_seeds(seed: int) -> None:
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-    random.seed(seed)
+def env_bool(name: str) -> bool:
+    return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"}
 
 
 def get_size_from_signature(signature: SignatureType | DetectionSignatureType) -> tuple[int, int]:

birder/common/training_cli.py

@@ -5,6 +5,7 @@ import typing
 from typing import Optional
 from typing import get_args
 
+from birder.common.cli import FlexibleDictAction
 from birder.common.cli import ValidationError
 from birder.common.training_utils import OptimizerType
 from birder.common.training_utils import SchedulerType
@@ -82,11 +83,23 @@ def add_lr_wd_args(parser: argparse.ArgumentParser, backbone_lr: bool = False, w
         metavar="WD",
         help="weight decay for embedding parameters for vision transformer models",
     )
+    group.add_argument(
+        "--custom-layer-wd",
+        action=FlexibleDictAction,
+        metavar="LAYER=WD",
+        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)")
     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"
     )
+    group.add_argument(
+        "--custom-layer-lr-scale",
+        action=FlexibleDictAction,
+        metavar="LAYER=SCALE",
+        help="custom lr_scale for specific layers by name (e.g., offset_conv=0.01,attention=0.5)",
+    )
 
 
 def add_lr_scheduler_args(parser: argparse.ArgumentParser) -> None:
@@ -110,10 +123,13 @@ def add_lr_scheduler_args(parser: argparse.ArgumentParser) -> None:
         type=int,
         default=40,
         metavar="N",
-        help="decrease lr every step-size epochs (for step scheduler only)",
+        help="decrease lr every N epochs/steps (relative to after warmup, step scheduler only)",
     )
     group.add_argument(
-        "--lr-steps", type=int, nargs="+", help="decrease lr every step-size epochs (multistep scheduler only)"
+        "--lr-steps",
+        type=int,
+        nargs="+",
+        help="absolute epoch/step milestones when to decrease lr (multistep scheduler only)",
     )
     group.add_argument(
         "--lr-step-gamma",
@@ -182,6 +198,11 @@ def add_detection_input_args(parser: argparse.ArgumentParser) -> None:
         action="store_true",
         help="enable random square resize once per batch (capped by max(--size))",
     )
+    group.add_argument(
+        "--multiscale-min-size",
+        type=int,
+        help="minimum short-edge size for multiscale lists (rounded up to nearest multiple of 32)",
+    )
 
 
 def add_training_schedule_args(parser: argparse.ArgumentParser, default_epochs: int = 100) -> None:
@@ -391,7 +412,7 @@ def add_ema_args(
         "--model-ema-warmup",
         type=int,
         metavar="N",
-        help="number of epochs before EMA is applied (defaults to warmup epochs/iters, pass 0 to disable warmup)",
+        help="number of epochs/steps before EMA is applied (defaults to warmup epochs/steps, pass 0 to disable warmup)",
     )
 
 
@@ -656,6 +677,11 @@ def common_args_validation(args: argparse.Namespace) -> None:
                 f"but it is set to '{args.lr_scheduler_update}'"
             )
 
+    # EMA
+    if hasattr(args, "model_ema_steps") is True:
+        if args.model_ema_steps < 1:
+            raise ValidationError("--model-ema-steps must be >= 1")
+
     # Compile args, argument dependant
     if hasattr(args, "compile_teacher") is True:
         if args.compile is True and args.compile_teacher is True:

birder/common/training_utils.py

@@ -3,8 +3,10 @@ import contextlib
 import logging
 import math
 import os
+import random
 import re
 import subprocess
+import sys
 from collections import deque
 from collections.abc import Callable
 from collections.abc import Generator
@@ -29,12 +31,25 @@ from birder.data.transforms.classification import training_preset
 from birder.optim import Lamb
 from birder.optim import Lars
 from birder.scheduler import CooldownLR
+from birder.version import __version__ as birder_version
 
 logger = logging.getLogger(__name__)
 
 OptimizerType = Literal["sgd", "rmsprop", "adam", "adamw", "nadam", "nadamw", "lamb", "lambw", "lars"]
 SchedulerType = Literal["constant", "step", "multistep", "cosine", "polynomial"]
 
+###############################################################################
+# Core Utilities
+###############################################################################
+
+
+def set_random_seeds(seed: int) -> None:
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
 ###############################################################################
 # Data Sampling
 ###############################################################################
@@ -207,13 +222,16 @@ def count_layers(model: torch.nn.Module) -> int:
 def optimizer_parameter_groups(
     model: torch.nn.Module,
     weight_decay: float,
+    base_lr: float,
     norm_weight_decay: Optional[float] = None,
     custom_keys_weight_decay: Optional[list[tuple[str, float]]] = None,
+    custom_layer_weight_decay: Optional[dict[str, float]] = None,
     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,
     backbone_lr: Optional[float] = None,
+    custom_layer_lr_scale: Optional[dict[str, float]] = None,
 ) -> list[dict[str, Any]]:
     """
     Return parameter groups for optimizers with per-parameter group weight decay.
@@ -233,11 +251,16 @@ def optimizer_parameter_groups(
         The PyTorch model whose parameters will be grouped for optimization.
     weight_decay
         Default weight decay (L2 regularization) value applied to parameters.
+    base_lr
+        Base learning rate that will be scaled by lr_scale factors for each parameter group.
     norm_weight_decay
         Weight decay value specifically for normalization layers. If None, uses weight_decay.
     custom_keys_weight_decay
         List of (parameter_name, weight_decay) tuples for applying custom weight decay
         values to specific parameters by name matching.
+    custom_layer_weight_decay
+        Dictionary mapping layer name substrings to custom weight decay values.
+        Applied to parameters whose names contain the specified keys.
     layer_decay
         Layer-wise learning rate decay factor.
     layer_decay_min_scale
@@ -248,6 +271,9 @@ def optimizer_parameter_groups(
         Custom learning rate for bias parameters (parameters ending with '.bias').
     backbone_lr
         Custom learning rate for backbone parameters (parameters starting with 'backbone.').
+    custom_layer_lr_scale
+        Dictionary mapping layer name substrings to custom lr_scale values.
+        Applied to parameters whose names contain the specified keys.
 
     Returns
     -------
@@ -291,14 +317,14 @@ def optimizer_parameter_groups(
     if layer_decay is not None:
         layer_max = num_layers - 1
         layer_scales = [max(layer_decay_min_scale, layer_decay ** (layer_max - i)) for i in range(num_layers)]
-        logger.info(f"Layer scaling in range of {min(layer_scales)} - {max(layer_scales)} on {num_layers} layers")
+        logger.info(f"Layer scaling ranges from {min(layer_scales)} to {max(layer_scales)} across {num_layers} layers")
 
     # Set weight decay and layer decay
     idx = 0
     params = []
     module_stack_with_prefix = [(model, "")]
     visited_modules = []
-    while len(module_stack_with_prefix) > 0:
+    while len(module_stack_with_prefix) > 0:  # pylint: disable=too-many-nested-blocks
         skip_module = False
         (module, prefix) = module_stack_with_prefix.pop()
         if id(module) in visited_modules:
@@ -324,13 +350,35 @@ def optimizer_parameter_groups(
                 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]
+                        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:
+                                    lr_scale = custom_scale
+                                    break
+
+                        # Apply custom layer weight decay (substring matching)
+                        wd = custom_wd
+                        if custom_layer_weight_decay is not None:
+                            for layer_name_key, custom_wd_value in custom_layer_weight_decay.items():
+                                if layer_name_key in target_name:
+                                    wd = custom_wd_value
+                                    break
+
                         d = {
                             "params": p,
-                            "weight_decay": custom_wd,
-                            "lr_scale": 1.0 if layer_decay is None else layer_scales[idx],
+                            "weight_decay": wd,
+                            "lr_scale": lr_scale,  # Used only for reference/debugging
                         }
-                        if backbone_lr is not None and target_name.startswith("backbone.") is True:
+
+                        # 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 lr_scale != 1.0:
+                            d["lr"] = base_lr * lr_scale
 
                         params.append(d)
                         is_custom_key = True
@@ -342,16 +390,34 @@ def optimizer_parameter_groups(
                 else:
                     wd = weight_decay
 
+                # Apply custom layer weight decay (substring matching)
+                if custom_layer_weight_decay is not None:
+                    for layer_name_key, custom_wd_value in custom_layer_weight_decay.items():
+                        if layer_name_key in target_name:
+                            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]
+                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:
+                            lr_scale = custom_scale
+                            break
+
                 d = {
                     "params": p,
                     "weight_decay": wd,
-                    "lr_scale": 1.0 if layer_decay is None else layer_scales[idx],
+                    "lr_scale": lr_scale,  # Used only for reference/debugging
                 }
-                if backbone_lr is not None and target_name.startswith("backbone.") is True:
-                    d["lr"] = backbone_lr
 
+                # 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 lr_scale != 1.0:
+                    d["lr"] = base_lr * lr_scale
 
                 params.append(d)
 
@@ -442,6 +508,8 @@ def get_optimizer(parameters: list[dict[str, Any]], l_rate: float, args: argpars
     else:
         raise ValueError("Unknown optimizer")
 
+    logger.debug(f"Created {opt} optimizer with lr={lr}, weight_decay={args.wd}")
+
     return optimizer
 
 
@@ -477,10 +545,10 @@ def get_scheduler(
 
     main_steps = steps - begin_step - remaining_warmup - remaining_cooldown - 1
 
-    logger.debug(f"Using {steps_per_epoch} steps per epoch")
+    logger.debug(f"Scheduler using {steps_per_epoch} steps per epoch")
     logger.debug(
         f"Scheduler {args.lr_scheduler} set for {steps} steps of which {warmup_steps} "
-        f"are warmup and {cooldown_steps} cooldown"
+        f"are warmup and {cooldown_steps} are cooldown"
     )
     logger.debug(
         f"Currently starting from step {begin_step} with {remaining_warmup} remaining warmup steps "
@@ -491,12 +559,29 @@ def get_scheduler(
     if args.lr_scheduler == "constant":
         main_scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=1)
     elif args.lr_scheduler == "step":
+        # Note: StepLR step_size is relative to when the main scheduler starts (after warmup)
+        # This means drops occur relative to the end of warmup, not at absolute epoch numbers
         main_scheduler = torch.optim.lr_scheduler.StepLR(
             optimizer, step_size=args.lr_step_size, gamma=args.lr_step_gamma
         )
     elif args.lr_scheduler == "multistep":
+        # For MultiStepLR, milestones should be absolute step numbers
+        # Adjust them to be relative to when the main scheduler starts (after warmup)
+        # This ensures drops occur at the specified absolute steps, not relative to after warmup
+        adjusted_milestones = [m - warmup_steps for m in args.lr_steps if m >= warmup_steps]
+        if len(adjusted_milestones) == 0:
+            logger.debug(
+                f"All MultiStepLR milestones {args.lr_steps} are before warmup "
+                f"(warmup ends at step {warmup_steps}). Using empty milestone list."
+            )
+            adjusted_milestones = []
+
+        logger.debug(
+            f"MultiStepLR milestones adjusted from {args.lr_steps} to {adjusted_milestones} "
+            f"(relative to main scheduler start after {warmup_steps} warmup steps)"
+        )
         main_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-            optimizer, milestones=args.lr_steps, gamma=args.lr_step_gamma
+            optimizer, milestones=adjusted_milestones, gamma=args.lr_step_gamma
         )
     elif args.lr_scheduler == "cosine":
         main_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
@@ -793,6 +878,51 @@ def is_local_primary(args: argparse.Namespace) -> bool:
     return args.local_rank == 0  # type: ignore[no-any-return]
 
 
+def init_training(
+    args: argparse.Namespace,
+    log: logging.Logger,
+    *,
+    cudnn_dynamic_size: bool = False,
+) -> tuple[torch.device, int, bool]:
+    init_distributed_mode(args)
+
+    log.info(f"Starting training, birder version: {birder_version}, pytorch version: {torch.__version__}")
+
+    log_git_info()
+
+    if args.cpu is True:
+        device = torch.device("cpu")
+        device_id = 0
+    else:
+        device = torch.device("cuda")
+        device_id = torch.cuda.current_device()
+
+    if args.use_deterministic_algorithms is True:
+        torch.backends.cudnn.benchmark = False
+        torch.use_deterministic_algorithms(True)
+    elif cudnn_dynamic_size is True:
+        # Dynamic sizes: avoid per-size algorithm selection overhead.
+        torch.backends.cudnn.enabled = False
+    else:
+        torch.backends.cudnn.enabled = True
+        torch.backends.cudnn.benchmark = True
+
+    if args.seed is not None:
+        set_random_seeds(args.seed)
+
+    if args.non_interactive is True or is_local_primary(args) is False:
+        disable_tqdm = True
+    elif sys.stderr.isatty() is False:
+        disable_tqdm = True
+    else:
+        disable_tqdm = False
+
+    # Enable or disable the autograd anomaly detection.
+    torch.autograd.set_detect_anomaly(args.grad_anomaly_detection)
+
+    return (device, device_id, disable_tqdm)
+
+
 ###############################################################################
 # Utility Functions
 ###############################################################################

birder/data/collators/detection.py

@@ -1,13 +1,14 @@
 import math
 import random
 from typing import Any
+from typing import Optional
 
 import torch
 from torchvision import tv_tensors
 from torchvision.transforms import v2
 from torchvision.transforms.v2 import functional as F
 
-BATCH_MULTISCALE_SIZES = (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
+from birder.data.transforms.detection import build_multiscale_sizes
 
 
 def collate_fn(batch: list[tuple[Any, ...]]) -> tuple[Any, ...]:
@@ -63,13 +64,19 @@ class DetectionCollator:
 
 
 class BatchRandomResizeCollator(DetectionCollator):
-    def __init__(self, input_offset: int, size: tuple[int, int], size_divisible: int = 32) -> None:
+    def __init__(
+        self,
+        input_offset: int,
+        size: tuple[int, int],
+        size_divisible: int = 32,
+        multiscale_min_size: Optional[int] = None,
+    ) -> None:
         super().__init__(input_offset, size_divisible=size_divisible)
         if size is None:
             raise ValueError("size must be provided for batch multiscale")
 
         max_side = max(size)
-        sizes = [side for side in BATCH_MULTISCALE_SIZES if side <= max_side]
+        sizes = [side for side in build_multiscale_sizes(multiscale_min_size) if side <= max_side]
         if len(sizes) == 0:
             sizes = [max_side]
 

birder/data/datasets/coco.py

@@ -98,10 +98,14 @@ class CocoTraining(CocoBase):
 class CocoInference(CocoBase):
     def __getitem__(self, index: int) -> tuple[str, torch.Tensor, Any, list[int]]:
         coco_id = self.dataset.ids[index]
-        path = self.dataset.coco.loadImgs(coco_id)[0]["file_name"]
+        img_info = self.dataset.coco.loadImgs(coco_id)[0]
+        path = img_info["file_name"]
         (sample, labels) = self.dataset[index]
 
-        return (path, sample, labels, F.get_size(sample))
+        # Get original image size (height, width) before transforms
+        orig_size = [img_info["height"], img_info["width"]]
+
+        return (path, sample, labels, orig_size)
 
 
 class CocoMosaicTraining(CocoBase):
@@ -127,9 +131,7 @@ class CocoMosaicTraining(CocoBase):
         self._mosaic_decay_epochs: Optional[int] = None
         self._mosaic_decay_start: Optional[int] = None
 
-    def configure_mosaic_linear_decay(
-        self, base_prob: float, total_epochs: int, decay_fraction: float = 0.1
-    ) -> None:
+    def configure_mosaic_linear_decay(self, base_prob: float, total_epochs: int, decay_fraction: float = 0.1) -> None:
         if total_epochs <= 0:
             raise ValueError("total_epochs must be positive")
         if decay_fraction <= 0.0 or decay_fraction > 1.0:
@@ -141,11 +143,7 @@ class CocoMosaicTraining(CocoBase):
         self._mosaic_decay_start = max(1, total_epochs - decay_epochs + 1)
 
     def update_mosaic_prob(self, epoch: int) -> Optional[float]:
-        if (
-            self._mosaic_base_prob is None
-            or self._mosaic_decay_epochs is None
-            or self._mosaic_decay_start is None
-        ):
+        if self._mosaic_base_prob is None or self._mosaic_decay_epochs is None or self._mosaic_decay_start is None:
             return None
 
         if epoch >= self._mosaic_decay_start:

birder/data/transforms/detection.py

@@ -1,3 +1,4 @@
+import math
 import random
 from collections.abc import Callable
 from typing import Any
@@ -10,6 +11,24 @@ from torchvision.transforms import v2
 
 from birder.data.transforms.classification import RGBType
 
+MULTISCALE_STEP = 32
+DEFAULT_MULTISCALE_MIN_SIZE = 480
+DEFAULT_MULTISCALE_MAX_SIZE = 800
+
+
+def build_multiscale_sizes(
+    min_size: Optional[int] = None, max_size: int = DEFAULT_MULTISCALE_MAX_SIZE
+) -> tuple[int, ...]:
+    if min_size is None:
+        min_size = DEFAULT_MULTISCALE_MIN_SIZE
+
+    start = int(math.ceil(min_size / MULTISCALE_STEP) * MULTISCALE_STEP)
+    end = int(math.floor(max_size / MULTISCALE_STEP) * MULTISCALE_STEP)
+    if end < start:
+        return (start,)
+
+    return tuple(range(start, end + 1, MULTISCALE_STEP))
+
 
 class ResizeWithRandomInterpolation(nn.Module):
     def __init__(
@@ -39,6 +58,7 @@ def get_birder_augment(
     dynamic_size: bool,
     multiscale: bool,
     max_size: Optional[int],
+    multiscale_min_size: Optional[int],
     post_mosaic: bool = False,
 ) -> Callable[..., torch.Tensor]:
     if dynamic_size is True:
@@ -78,9 +98,7 @@ def get_birder_augment(
     # Resize
     if multiscale is True:
         transformations.append(
-            v2.RandomShortestSize(
-                min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=max_size or 1333
-            ),
+            v2.RandomShortestSize(min_size=build_multiscale_sizes(multiscale_min_size), max_size=max_size or 1333),
         )
     else:
         transformations.append(
@@ -132,6 +150,7 @@ def get_birder_augment(
 AugType = Literal["birder", "lsj", "multiscale", "ssd", "ssdlite", "yolo", "detr"]
 
 
+# pylint: disable=too-many-return-statements
 def training_preset(
     size: tuple[int, int],
     aug_type: AugType,
@@ -140,6 +159,7 @@ def training_preset(
     dynamic_size: bool = False,
     multiscale: bool = False,
     max_size: Optional[int] = None,
+    multiscale_min_size: Optional[int] = None,
     post_mosaic: bool = False,
 ) -> Callable[..., torch.Tensor]:
     mean = rgv_values["mean"]
@@ -159,7 +179,9 @@ def training_preset(
         return v2.Compose(  # type:ignore
             [
                 v2.ToImage(),
-                get_birder_augment(size, level, fill_value, dynamic_size, multiscale, max_size, post_mosaic),
+                get_birder_augment(
+                    size, level, fill_value, dynamic_size, multiscale, max_size, multiscale_min_size, post_mosaic
+                ),
                 v2.ToDtype(torch.float32, scale=True),
                 v2.Normalize(mean=mean, std=std),
                 v2.ToPureTensor(),
@@ -190,9 +212,7 @@ def training_preset(
         return v2.Compose(  # type: ignore
             [
                 v2.ToImage(),
-                v2.RandomShortestSize(
-                    min_size=(480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=max_size or 1333
-                ),
+                v2.RandomShortestSize(min_size=build_multiscale_sizes(multiscale_min_size), max_size=max_size or 1333),
                 v2.RandomHorizontalFlip(0.5),
                 v2.SanitizeBoundingBoxes(),
                 v2.ToDtype(torch.float32, scale=True),
@@ -264,21 +284,18 @@ def training_preset(
         )
 
     if aug_type == "detr":
+        multiscale_sizes = build_multiscale_sizes(multiscale_min_size)
         return v2.Compose(  # type: ignore
             [
                 v2.ToImage(),
                 v2.RandomChoice(
                     [
-                        v2.RandomShortestSize(
-                            (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=max_size or 1333
-                        ),
+                        v2.RandomShortestSize(min_size=multiscale_sizes, max_size=max_size or 1333),
                         v2.Compose(
                             [
                                 v2.RandomShortestSize((400, 500, 600)),
                                 v2.RandomIoUCrop() if post_mosaic is False else v2.Identity(),  # RandomSizeCrop
-                                v2.RandomShortestSize(
-                                    (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800), max_size=max_size or 1333
-                                ),
+                                v2.RandomShortestSize(min_size=multiscale_sizes, max_size=max_size or 1333),
                             ]
                         ),
                     ]

birder/inference/data_parallel.py

@@ -1,8 +1,7 @@
 """
 Inference-optimized multi-GPU parallelization
 
-This module provides InferenceDataParallel, an inference-specific alternative to
-torch.nn.DataParallel.
+This module provides InferenceDataParallel, an inference-specific alternative to torch.nn.DataParallel.
 """
 
 import copy