birder 0.2.2__tar.gz → 0.3.0__tar.gz

{birder-0.2.2 → birder-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: birder
-Version: 0.2.2
+Version: 0.3.0
 Summary: An open-source computer vision framework for wildlife image analysis, featuring state-of-the-art models for species classification and detection.
 Author: Ofer Hasson
 License-Expression: Apache-2.0
@@ -45,7 +45,7 @@ Provides-Extra: dev
 Requires-Dist: altair~=5.5.0; extra == "dev"
 Requires-Dist: bandit~=1.9.2; extra == "dev"
 Requires-Dist: black~=25.12.0; extra == "dev"
-Requires-Dist: build~=1.3.0; extra == "dev"
+Requires-Dist: build~=1.4.0; extra == "dev"
 Requires-Dist: bumpver~=2025.1131; extra == "dev"
 Requires-Dist: captum~=0.7.0; extra == "dev"
 Requires-Dist: coverage~=7.13.1; extra == "dev"
@@ -66,6 +66,7 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: requests~=2.32.5; extra == "dev"
 Requires-Dist: safetensors~=0.7.0; extra == "dev"
 Requires-Dist: setuptools; extra == "dev"
+Requires-Dist: torchao~=0.15.0; extra == "dev"
 Requires-Dist: torchprofile==0.0.4; extra == "dev"
 Requires-Dist: twine~=6.2.0; extra == "dev"
 Requires-Dist: types-requests~=2.32.4; extra == "dev"
@@ -208,7 +209,7 @@ For detailed information about these datasets, including descriptions, citations
 
 ## Detection
 
-Detection training and inference are available, see [docs/training_detection.md](docs/training_detection.md) and
+Detection training and inference are available, see [docs/training_scripts.md](docs/training_scripts.md) and
 [docs/inference.md](docs/inference.md). APIs and model coverage may evolve as detection support matures.
 
 ## Project Status and Contributions

{birder-0.2.2 → birder-0.3.0}/README.md

@@ -129,7 +129,7 @@ For detailed information about these datasets, including descriptions, citations
 
 ## Detection
 
-Detection training and inference are available, see [docs/training_detection.md](docs/training_detection.md) and
+Detection training and inference are available, see [docs/training_scripts.md](docs/training_scripts.md) and
 [docs/inference.md](docs/inference.md). APIs and model coverage may evolve as detection support matures.
 
 ## Project Status and Contributions

{birder-0.2.2 → birder-0.3.0}/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-0.2.2 → birder-0.3.0}/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:
@@ -185,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:
@@ -193,6 +211,12 @@ def add_training_schedule_args(parser: argparse.ArgumentParser, default_epochs:
     group.add_argument(
         "--stop-epoch", type=int, metavar="N", help="epoch to stop the training at (multi stage training)"
     )
+    group.add_argument(
+        "--steps-per-epoch",
+        type=int,
+        metavar="N",
+        help="virtual epoch length in steps, leave unset to use the full dataset",
+    )
     group.add_argument("--warmup-epochs", type=int, metavar="N", help="number of warmup epochs")
     group.add_argument("--warmup-steps", type=int, metavar="N", help="number of warmup optimizer steps")
     group.add_argument("--cooldown-epochs", type=int, metavar="N", help="number of cooldown epochs (linear to zero)")

{birder-0.2.2 → birder-0.3.0}/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
@@ -15,6 +17,7 @@ from typing import Any
 from typing import Literal
 from typing import Optional
 from typing import Sized
+from typing import overload
 
 import numpy as np
 import torch
@@ -29,12 +32,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
 ###############################################################################
@@ -55,13 +71,7 @@ class RASampler(torch.utils.data.Sampler):
     """
 
     def __init__(
-        self,
-        dataset: Sized,
-        num_replicas: int,
-        rank: int,
-        shuffle: bool,
-        seed: int = 0,
-        repetitions: int = 3,
+        self, dataset: Sized, num_replicas: int, rank: int, shuffle: bool, seed: int = 0, repetitions: int = 3
     ) -> None:
         super().__init__()
         self.dataset = dataset
@@ -70,12 +80,11 @@ class RASampler(torch.utils.data.Sampler):
         self.epoch = 0
         self.num_samples = int(math.ceil(len(self.dataset) * float(repetitions) / self.num_replicas))
         self.total_size = self.num_samples * self.num_replicas
-        self.num_selected_samples = int(math.floor(len(self.dataset) // 256 * 256 / self.num_replicas))
         self.shuffle = shuffle
         self.seed = seed
         self.repetitions = repetitions
 
-    def __iter__(self) -> Iterator[list[int]]:
+    def __iter__(self) -> Iterator[int]:
         if self.shuffle is True:
             # Deterministically shuffle based on epoch
             g = torch.Generator()
@@ -85,18 +94,148 @@ class RASampler(torch.utils.data.Sampler):
             indices = list(range(len(self.dataset)))
 
         # Add extra samples to make it evenly divisible
-        indices = [ele for ele in indices for i in range(self.repetitions)]
-        indices += indices[: (self.total_size - len(indices))]
-        assert len(indices) == self.total_size
+        indices = [ele for ele in indices for _ in range(self.repetitions)]
+        if len(indices) < self.total_size:
+            indices += indices[: (self.total_size - len(indices))]
+        else:
+            indices = indices[: self.total_size]
 
-        # Subsample
+        # Shard by rank
         indices = indices[self.rank : self.total_size : self.num_replicas]
         assert len(indices) == self.num_samples
 
-        return iter(indices[: self.num_selected_samples])
+        yield from indices
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+
+class InfiniteSampler(torch.utils.data.Sampler):
+    """
+    Infinite sampler that loops indefinitely over the dataset
+    """
+
+    def __init__(self, dataset: Sized, shuffle: bool, seed: int = 0) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.shuffle = shuffle
+        self.seed = seed
+        self.epoch = 0
+
+    def __iter__(self) -> Iterator[int]:
+        g = torch.Generator()
+        while True:
+            if self.shuffle is True:
+                g.manual_seed(self.seed + self.epoch)
+                indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            else:
+                indices = list(range(len(self.dataset)))
+
+            yield from indices
+
+            logger.debug(f"InfiniteSampler finished epoch {self.epoch}")
+            self.epoch += 1
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+
+class InfiniteDistributedSampler(torch.utils.data.Sampler):
+    """
+    Infinite distributed sampler that keeps a continuous shuffled stream per rank
+    """
+
+    def __init__(self, dataset: Sized, num_replicas: int, rank: int, shuffle: bool, seed: int = 0) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.shuffle = shuffle
+        self.seed = seed
+        self.epoch = 0
+        self.num_samples = int(math.ceil(len(self.dataset) / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self) -> Iterator[int]:
+        g = torch.Generator()
+        while True:
+            if self.shuffle is True:
+                g.manual_seed(self.seed + self.epoch)
+                indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            else:
+                indices = list(range(len(self.dataset)))
+
+            if len(indices) < self.total_size:
+                indices += indices[: (self.total_size - len(indices))]
+            else:
+                indices = indices[: self.total_size]
+
+            indices = indices[self.rank : self.total_size : self.num_replicas]
+            assert len(indices) == self.num_samples
+
+            yield from indices
+
+            logger.debug(f"InfiniteDistributedSampler finished epoch {self.epoch}")
+            self.epoch += 1
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+
+class InfiniteRASampler(torch.utils.data.Sampler):
+    """
+    Infinite version of the repeated augmentation sampler
+    """
+
+    def __init__(
+        self, dataset: Sized, num_replicas: int, rank: int, shuffle: bool, seed: int = 0, repetitions: int = 3
+    ) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.num_samples = int(math.ceil(len(self.dataset) * float(repetitions) / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+        self.shuffle = shuffle
+        self.seed = seed
+        self.repetitions = repetitions
+
+    def __iter__(self) -> Iterator[int]:
+        g = torch.Generator()
+        while True:
+            if self.shuffle is True:
+                g.manual_seed(self.seed + self.epoch)
+                indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            else:
+                indices = list(range(len(self.dataset)))
+
+            indices = [ele for ele in indices for _ in range(self.repetitions)]
+            if len(indices) < self.total_size:
+                indices += indices[: (self.total_size - len(indices))]
+            else:
+                indices = indices[: self.total_size]
+
+            # Shard by rank
+            indices = indices[self.rank : self.total_size : self.num_replicas]
+            assert len(indices) == self.num_samples
+
+            yield from indices
+
+            logger.debug(f"InfiniteRASampler finished epoch {self.epoch}")
+            self.epoch += 1
 
     def __len__(self) -> int:
-        return self.num_selected_samples
+        return self.num_samples
 
     def set_epoch(self, epoch: int) -> None:
         self.epoch = epoch
@@ -207,13 +346,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 +375,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 +395,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 +441,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 +474,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 +514,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 +632,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 +669,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 "
@@ -568,27 +760,87 @@ def get_amp_scaler(amp: bool, amp_dtype_str: str) -> tuple[Optional[torch.amp.Gr
     return (scaler, amp_dtype)
 
 
+@overload
 def get_samplers(
-    args: argparse.Namespace, training_dataset: torch.utils.data.Dataset, validation_dataset: torch.utils.data.Dataset
-) -> torch.utils.data.Sampler:
-    if args.distributed is True:
-        if args.ra_sampler is True:
-            train_sampler = RASampler(
-                training_dataset,
-                num_replicas=args.world_size,
-                rank=args.rank,
-                shuffle=True,
-                repetitions=args.ra_reps,
-            )
+    args: argparse.Namespace,
+    training_dataset: torch.utils.data.Dataset,
+    validation_dataset: torch.utils.data.Dataset,
+    infinite: bool = False,
+) -> tuple[torch.utils.data.Sampler, torch.utils.data.Sampler]: ...
 
-        else:
-            train_sampler = torch.utils.data.distributed.DistributedSampler(training_dataset, shuffle=True)
 
-        validation_sampler = torch.utils.data.distributed.DistributedSampler(validation_dataset, shuffle=False)
+@overload
+def get_samplers(
+    args: argparse.Namespace,
+    training_dataset: torch.utils.data.Dataset,
+    validation_dataset: None = None,
+    infinite: bool = False,
+) -> tuple[torch.utils.data.Sampler, None]: ...
+
+
+def get_samplers(
+    args: argparse.Namespace,
+    training_dataset: torch.utils.data.Dataset,
+    validation_dataset: Optional[torch.utils.data.Dataset] = None,
+    infinite: bool = False,
+) -> tuple[torch.utils.data.Sampler, Optional[torch.utils.data.Sampler]]:
+    if args.seed is None:
+        seed = int(torch.empty((), dtype=torch.int64).random_().item())
+        if is_dist_available_and_initialized() is True:
+            seed_tensor = torch.tensor(seed, dtype=torch.int64).cuda()
+            dist.broadcast(seed_tensor, src=0, async_op=False)
+            seed = int(seed_tensor.item())
+    else:
+        seed = args.seed
+
+    ra_sampler = getattr(args, "ra_sampler", False)
+    if args.distributed is True:
+        if infinite is True:
+            if ra_sampler is True:
+                train_sampler = InfiniteRASampler(
+                    training_dataset,
+                    num_replicas=args.world_size,
+                    rank=args.rank,
+                    shuffle=True,
+                    seed=seed,
+                    repetitions=args.ra_reps,
+                )
+            else:
+                train_sampler = InfiniteDistributedSampler(
+                    training_dataset, num_replicas=args.world_size, rank=args.rank, shuffle=True, seed=seed
+                )
+        else:
+            if ra_sampler is True:
+                train_sampler = RASampler(
+                    training_dataset,
+                    num_replicas=args.world_size,
+                    rank=args.rank,
+                    shuffle=True,
+                    seed=seed,
+                    repetitions=args.ra_reps,
+                )
+            else:
+                train_sampler = torch.utils.data.distributed.DistributedSampler(
+                    training_dataset, shuffle=True, seed=seed
+                )
+
+        if validation_dataset is None:
+            validation_sampler = None
+        else:
+            validation_sampler = torch.utils.data.distributed.DistributedSampler(validation_dataset, shuffle=False)
 
     else:
-        train_sampler = torch.utils.data.RandomSampler(training_dataset)
-        validation_sampler = torch.utils.data.SequentialSampler(validation_dataset)
+        if infinite is True:
+            train_sampler = InfiniteSampler(training_dataset, shuffle=True, seed=seed)
+        else:
+            generator = torch.Generator()
+            generator.manual_seed(seed)
+            train_sampler = torch.utils.data.RandomSampler(training_dataset, generator=generator)
+
+        if validation_dataset is None:
+            validation_sampler = None
+        else:
+            validation_sampler = torch.utils.data.SequentialSampler(validation_dataset)
 
     return (train_sampler, validation_sampler)
 
@@ -810,6 +1062,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
 ###############################################################################