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

birder/__init__.py

@@ -1,5 +1,6 @@
 from birder.common.fs_ops import load_model_with_cfg
 from birder.common.fs_ops import load_pretrained_model
+from birder.common.fs_ops import load_pretrained_model_and_transform
 from birder.common.lib import get_channels_from_signature
 from birder.common.lib import get_size_from_signature
 from birder.data.transforms.classification import inference_preset as classification_transform
@@ -17,5 +18,6 @@ __all__ = [
     "list_pretrained_models",
     "load_model_with_cfg",
     "load_pretrained_model",
+    "load_pretrained_model_and_transform",
     "__version__",
 ]

birder/common/fs_ops.py

@@ -2,6 +2,7 @@ import json
 import logging
 import os
 import re
+from collections.abc import Callable
 from collections.abc import Iterator
 from pathlib import Path
 from typing import Any
@@ -24,6 +25,8 @@ from birder.common.lib import get_network_name
 from birder.common.lib import get_pretrained_model_url
 from birder.conf import settings
 from birder.data.transforms.classification import RGBType
+from birder.data.transforms.classification import inference_preset
+from birder.data.transforms.detection import InferenceTransform
 from birder.model_registry import Task
 from birder.model_registry import registry
 from birder.model_registry.manifest import FileFormatType
@@ -801,7 +804,8 @@ def load_detection_model(
         for param in net.parameters():
             param.requires_grad_(False)
 
-        net.eval()
+        if pt2 is False:  # NOTE: Remove when GraphModule add support for 'eval'
+            net.eval()
 
     if len(backbone_loaded_config) == 0:
         backbone_custom_config = None
@@ -918,6 +922,82 @@ def load_pretrained_model(
     raise ValueError(f"Unknown model type: {model_metadata['task']}")
 
 
+def load_pretrained_model_and_transform(
+    weights: str,
+    *,
+    dst: Optional[str | Path] = None,
+    file_format: FileFormatType = "pt",
+    inference: bool = True,
+    device: Optional[torch.device] = None,
+    dtype: Optional[torch.dtype] = None,
+    custom_config: Optional[dict[str, Any]] = None,
+    progress_bar: bool = True,
+    classification_kwargs: Optional[dict[str, Any]] = None,
+    detection_kwargs: Optional[dict[str, Any]] = None,
+) -> tuple[BaseNet | DetectionBaseNet, ModelInfo | DetectionModelInfo, Callable[..., torch.Tensor]]:
+    """
+    Loads a pre-trained model and builds the matching inference transform
+
+    This is a convenience helper for the common inference path where the model and
+    its default preprocessing are needed together. Classification models use
+    inference_preset, detection models use InferenceTransform.
+
+    Parameters
+    ----------
+    weights
+        Name of the pre-trained weights to load from the model registry.
+    dst
+        Destination path where the model weights will be downloaded or loaded from.
+    file_format
+        Model format (e.g. pt, pt2, safetensors, etc.)
+    inference
+        Flag to prepare the model for inference mode.
+    device
+        The device to load the model on (cpu/cuda).
+    dtype
+        Data type for model parameters and computations (e.g., torch.float32, torch.float16).
+    custom_config
+        Additional model configuration that overrides or extends the predefined configuration.
+    progress_bar
+        Whether to display a progress bar during file download.
+    classification_kwargs
+        Optional keyword arguments forwarded to inference_preset.
+    detection_kwargs
+        Optional keyword arguments forwarded to InferenceTransform. If dynamic_size is
+        not provided it defaults to the model signature value.
+
+    Returns
+    -------
+    A tuple containing three elements:
+    - A PyTorch module (neural network model) loaded with pre-trained weights.
+    - Model info containing class mappings, signature, and RGB stats.
+    - An inference transform matching the model task.
+    """
+
+    net, model_info = load_pretrained_model(
+        weights,
+        dst=dst,
+        file_format=file_format,
+        inference=inference,
+        device=device,
+        dtype=dtype,
+        custom_config=custom_config,
+        progress_bar=progress_bar,
+    )
+
+    size = lib.get_size_from_signature(model_info.signature)
+    transform: Callable[..., torch.Tensor]
+    if isinstance(model_info, DetectionModelInfo):
+        detection_args = {} if detection_kwargs is None else dict(detection_kwargs)
+        detection_args.setdefault("dynamic_size", model_info.signature["dynamic"])
+        transform = InferenceTransform(size, model_info.rgb_stats, **detection_args)
+    else:
+        classification_args = {} if classification_kwargs is None else dict(classification_kwargs)
+        transform = inference_preset(size, model_info.rgb_stats, **classification_args)
+
+    return (net, model_info, transform)
+
+
 def load_model_with_cfg(
     cfg: dict[str, Any] | str | Path, weights_path: Optional[str | Path]
 ) -> tuple[torch.nn.Module, dict[str, Any]]:

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"
@@ -480,8 +485,13 @@ def add_dataloader_args(
         )
 
 
-def add_precision_args(parser: argparse.ArgumentParser) -> None:
+def add_precision_args(parser: argparse.ArgumentParser, channels_last: bool = False) -> None:
     group = parser.add_argument_group("Precision parameters")
+    if channels_last is True:
+        group.add_argument(
+            "--channels-last", default=False, action="store_true", help="use channels-last memory format"
+        )
+
     group.add_argument(
         "--model-dtype",
         type=str,

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
@@ -1108,12 +1165,16 @@ def init_training(
         device_id = torch.cuda.current_device()
 
     if args.use_deterministic_algorithms is True:
+        log.debug("Turning on deterministic algorithms")
         torch.backends.cudnn.benchmark = False
         torch.use_deterministic_algorithms(True)
     elif cudnn_dynamic_size is True:
         # Dynamic sizes: avoid per-size algorithm selection overhead.
+        log.debug("Turning off cudnn")
         torch.backends.cudnn.enabled = False
+        torch.backends.cudnn.benchmark = False
     else:
+        log.debug("Turning on cudnn")
         torch.backends.cudnn.enabled = True
         torch.backends.cudnn.benchmark = True
 

birder/data/collators/detection.py

@@ -15,7 +15,9 @@ def collate_fn(batch: list[tuple[Any, ...]]) -> tuple[Any, ...]:
     return tuple(zip(*batch))
 
 
-def batch_images(images: list[torch.Tensor], size_divisible: int) -> tuple[torch.Tensor, torch.Tensor, list[list[int]]]:
+def batch_images(
+    images: list[torch.Tensor], size_divisible: int
+) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int]]]:
     """
     Batch list of image tensors of different sizes into a single batch.
     Pad with zeros all images to the shape of the largest image in the list.

birder/datahub/_lib.py

@@ -1,5 +1,6 @@
 import logging
 import tarfile
+import zipfile
 from pathlib import Path
 
 from birder.common import cli
@@ -26,9 +27,17 @@ def download_url(url: str, target: str | Path, sha256: str, progress_bar: bool =
 
 def extract_archive(from_path: str | Path, to_path: str | Path) -> None:
     logger.info(f"Extracting {from_path} to {to_path}")
-    with tarfile.open(from_path, "r") as tar:
-        if hasattr(tarfile, "data_filter") is True:
-            tar.extractall(to_path, filter="data")
-        else:
-            # NOTE: Remove once minimum Python version is 3.12 or above
-            tar.extractall(to_path)  # nosec # tarfile_unsafe_members
+    if isinstance(from_path, str):
+        from_path = Path(from_path)
+
+    if from_path.suffix == ".zip":
+        with zipfile.ZipFile(from_path, "r") as zf:
+            zf.extractall(to_path)  # nosec # tarfile_unsafe_members
+
+    else:
+        with tarfile.open(from_path, "r") as tar:
+            if hasattr(tarfile, "data_filter") is True:
+                tar.extractall(to_path, filter="data")
+            else:
+                # NOTE: Remove once minimum Python version is 3.12 or above
+                tar.extractall(to_path)  # nosec # tarfile_unsafe_members