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

birder/scripts/train_simclr.py

@@ -33,7 +33,6 @@ from birder.common import training_utils
 from birder.common.lib import format_duration
 from birder.common.lib import get_mim_network_name
 from birder.common.lib import get_network_name
-from birder.common.lib import set_random_seeds
 from birder.conf import settings
 from birder.data.dataloader.webdataset import make_wds_loader
 from birder.data.datasets.directory import make_image_dataset
@@ -67,41 +66,13 @@ def train(args: argparse.Namespace) -> None:
     #
     # Initialize
     #
-    training_utils.init_distributed_mode(args)
-    logger.info(f"Starting training, birder version: {birder.__version__}, pytorch version: {torch.__version__}")
-    training_utils.log_git_info()
+    (device, device_id, disable_tqdm) = training_utils.init_training(args, logger)
 
     if args.size is None:
         args.size = registry.get_default_size(args.network)
 
     logger.info(f"Using size={args.size}")
 
-    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)
-    else:
-        torch.backends.cudnn.benchmark = True
-
-    if args.seed is not None:
-        set_random_seeds(args.seed)
-
-    if args.non_interactive is True or training_utils.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)
-
     #
     # Data
     #
@@ -148,7 +119,7 @@ def train(args: argparse.Namespace) -> None:
 
     batch_size: int = args.batch_size
     grad_accum_steps: int = args.grad_accum_steps
-    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
+    logger.debug(f"Effective batch size = {batch_size * grad_accum_steps * args.world_size}")
 
     # Data loaders and samplers
     if args.distributed is True:
@@ -189,6 +160,8 @@ def train(args: argparse.Namespace) -> None:
     else:
         args.stop_epoch += 1
 
+    logging.debug(f"Epoch has {last_batch_idx+1} iterations ({optimizer_steps_per_epoch} steps)")
+
     #
     # Initialize network
     #
@@ -236,22 +209,25 @@ def train(args: argparse.Namespace) -> None:
     # Optimizer, learning rate scheduler and training parameter groups
     #
 
+    # Learning rate scaling
+    lr = training_utils.scale_lr(args)
+
     # Training parameter groups
     custom_keys_weight_decay = training_utils.get_wd_custom_keys(args)
     parameters = training_utils.optimizer_parameter_groups(
         net,
         args.wd,
+        base_lr=lr,
         norm_weight_decay=args.norm_wd,
         custom_keys_weight_decay=custom_keys_weight_decay,
+        custom_layer_weight_decay=args.custom_layer_wd,
         layer_decay=args.layer_decay,
         layer_decay_min_scale=args.layer_decay_min_scale,
         layer_decay_no_opt_scale=args.layer_decay_no_opt_scale,
         bias_lr=args.bias_lr,
+        custom_layer_lr_scale=args.custom_layer_lr_scale,
     )
 
-    # Learning rate scaling
-    lr = training_utils.scale_lr(args)
-
     if args.lr_scheduler_update == "epoch":
         step_update = False
         scheduler_steps_per_epoch = 1

birder/scripts/train_vicreg.py

@@ -36,7 +36,6 @@ from birder.common import training_utils
 from birder.common.lib import format_duration
 from birder.common.lib import get_mim_network_name
 from birder.common.lib import get_network_name
-from birder.common.lib import set_random_seeds
 from birder.conf import settings
 from birder.data.dataloader.webdataset import make_wds_loader
 from birder.data.datasets.directory import make_image_dataset
@@ -70,41 +69,13 @@ def train(args: argparse.Namespace) -> None:
     #
     # Initialize
     #
-    training_utils.init_distributed_mode(args)
-    logger.info(f"Starting training, birder version: {birder.__version__}, pytorch version: {torch.__version__}")
-    training_utils.log_git_info()
+    (device, device_id, disable_tqdm) = training_utils.init_training(args, logger)
 
     if args.size is None:
         args.size = registry.get_default_size(args.network)
 
     logger.info(f"Using size={args.size}")
 
-    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)
-    else:
-        torch.backends.cudnn.benchmark = True
-
-    if args.seed is not None:
-        set_random_seeds(args.seed)
-
-    if args.non_interactive is True or training_utils.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)
-
     #
     # Data
     #
@@ -151,7 +122,7 @@ def train(args: argparse.Namespace) -> None:
 
     batch_size: int = args.batch_size
     grad_accum_steps: int = args.grad_accum_steps
-    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
+    logger.debug(f"Effective batch size = {batch_size * grad_accum_steps * args.world_size}")
 
     # Data loaders and samplers
     if args.distributed is True:
@@ -192,6 +163,8 @@ def train(args: argparse.Namespace) -> None:
     else:
         args.stop_epoch += 1
 
+    logging.debug(f"Epoch has {last_batch_idx+1} iterations ({optimizer_steps_per_epoch} steps)")
+
     #
     # Initialize network
     #
@@ -242,22 +215,25 @@ def train(args: argparse.Namespace) -> None:
     # Loss criteria, optimizer, learning rate scheduler and training parameter groups
     #
 
+    # Learning rate scaling
+    lr = training_utils.scale_lr(args)
+
     # Training parameter groups
     custom_keys_weight_decay = training_utils.get_wd_custom_keys(args)
     parameters = training_utils.optimizer_parameter_groups(
         net,
         args.wd,
+        base_lr=lr,
         norm_weight_decay=args.norm_wd,
         custom_keys_weight_decay=custom_keys_weight_decay,
+        custom_layer_weight_decay=args.custom_layer_wd,
         layer_decay=args.layer_decay,
         layer_decay_min_scale=args.layer_decay_min_scale,
         layer_decay_no_opt_scale=args.layer_decay_no_opt_scale,
         bias_lr=args.bias_lr,
+        custom_layer_lr_scale=args.custom_layer_lr_scale,
     )
 
-    # Learning rate scaling
-    lr = training_utils.scale_lr(args)
-
     if args.lr_scheduler_update == "epoch":
         step_update = False
         scheduler_steps_per_epoch = 1

birder/tools/auto_anchors.py

@@ -242,6 +242,7 @@ def _validate_args(
     return (size, num_scales, num_anchors, output_format, strides)
 
 
+# pylint: disable=too-many-locals
 def auto_anchors(args: argparse.Namespace) -> None:
     (size, num_scales, num_anchors, output_format, strides) = _validate_args(args)
 
@@ -272,6 +273,7 @@ def auto_anchors(args: argparse.Namespace) -> None:
     logger.info(f"Mean IoU: {best_iou.mean().item():.4f}")
 
     formatted_groups = _format_anchor_groups(anchor_groups, args.precision)
+    anchors_output = None
     if output_format == "pixels":
         if num_scales == 1:
             formatted_anchors: Any = formatted_groups[0]
@@ -280,6 +282,7 @@ def auto_anchors(args: argparse.Namespace) -> None:
 
         print("Anchors (pixels):")
         print(pformat(formatted_anchors))
+        anchors_output = formatted_anchors
 
     if output_format == "grid":
         grid_groups: list[torch.Tensor] = []
@@ -297,6 +300,21 @@ def auto_anchors(args: argparse.Namespace) -> None:
 
         print("Anchors (grid units):")
         print(pformat(formatted_grid_output))
+        anchors_output = formatted_grid_output
+
+    if args.output is not None:
+        payload = {
+            "anchors": anchors_output,
+            "format": output_format,
+            "size": [size[0], size[1]],
+        }
+        if output_format == "grid":
+            payload["strides"] = strides
+
+        with open(args.output, "w", encoding="utf-8") as handle:
+            json.dump(payload, handle, indent=2)
+
+        logger.info(f"Wrote anchors to {args.output}")
 
 
 def set_parser(subparsers: Any) -> None:
@@ -312,7 +330,7 @@ def set_parser(subparsers: Any) -> None:
             "python -m birder.tools auto-anchors --size 640 --num-anchors 9 --num-scales 3 --format pixels "
             "--coco-json-path data/detection_data/training_annotations_coco.json\n"
             "python -m birder.tools auto-anchors --preset yolo_v4_tiny --size 416 416 "
-            "--coco-json-path ~/Datasets/cocodataset/annotations/instances_train2017.json\n"
+            "--coco-json-path ~/Datasets/cocodataset/annotations/instances_train2017.json --output anchors.json\n"
             "python -m birder.tools auto-anchors --preset yolo_v2 --stride 32 "
             "--coco-json-path data/detection_data/training_annotations_coco.json\n"
             "python -m birder.tools auto-anchors --size 640 --num-anchors 9 --num-scales 3 "
@@ -354,6 +372,7 @@ def set_parser(subparsers: Any) -> None:
         default=f"{settings.TRAINING_DETECTION_ANNOTATIONS_PATH}_coco.json",
         help="training COCO json path",
     )
+    subparser.add_argument("--output", type=str, help="write anchors as JSON to this path")
     subparser.set_defaults(func=main)
 
 

birder/tools/pack.py

@@ -3,6 +3,7 @@ import json
 import logging
 import multiprocessing
 import os
+import queue
 import signal
 import time
 from collections.abc import Callable
@@ -67,39 +68,43 @@ def _save_classes(pack_path: Path, class_to_idx: dict[str, int]) -> None:
 
 
 def _encode_image(path: str, file_format: str, size: Optional[int] = None) -> bytes:
-    image: Image.Image = Image.open(path)
-    if size is not None and size < min(image.size):
-        if image.size[0] > image.size[1]:
-            ratio = image.size[1] / size
-        else:
-            ratio = image.size[0] / size
-
-        width = round(image.size[0] / ratio)
-        height = round(image.size[1] / ratio)
-        if max(width, height) > MAX_SIZE:
-            if width > height:
-                ratio = width / MAX_SIZE
+    image: Image.Image
+    with Image.open(path) as image:
+        if file_format.lower() in ("jpeg", "jpg") and image.mode in ("RGBA", "P"):
+            image = image.convert("RGB")
+
+        if size is not None and size < min(image.size):
+            if image.size[0] > image.size[1]:
+                ratio = image.size[1] / size
             else:
-                ratio = height / MAX_SIZE
+                ratio = image.size[0] / size
+
+            width = round(image.size[0] / ratio)
+            height = round(image.size[1] / ratio)
+            if max(width, height) > MAX_SIZE:
+                if width > height:
+                    ratio = width / MAX_SIZE
+                else:
+                    ratio = height / MAX_SIZE
 
-            width = round(width / ratio)
-            height = round(height / ratio)
+                width = round(width / ratio)
+                height = round(height / ratio)
 
-        image = image.resize((width, height), Image.Resampling.BICUBIC)
+            image = image.resize((width, height), Image.Resampling.BICUBIC)
 
-    elif max(image.size) > MAX_SIZE:
-        if image.size[0] > image.size[1]:
-            ratio = image.size[0] / MAX_SIZE
-        else:
-            ratio = image.size[1] / MAX_SIZE
+        elif max(image.size) > MAX_SIZE:
+            if image.size[0] > image.size[1]:
+                ratio = image.size[0] / MAX_SIZE
+            else:
+                ratio = image.size[1] / MAX_SIZE
 
-        width = round(image.size[0] / ratio)
-        height = round(image.size[1] / ratio)
-        image = image.resize((width, height), Image.Resampling.BICUBIC)
+            width = round(image.size[0] / ratio)
+            height = round(image.size[1] / ratio)
+            image = image.resize((width, height), Image.Resampling.BICUBIC)
 
-    sample_buffer = BytesIO()
-    image.save(sample_buffer, format=file_format, quality=85)
-    return sample_buffer.getvalue()
+        sample_buffer = BytesIO()
+        image.save(sample_buffer, format=file_format, quality=85)
+        return sample_buffer.getvalue()
 
 
 def read_worker(q_in: Any, q_out: Any, error_event: Any, size: Optional[int], file_format: str) -> None:
@@ -162,38 +167,43 @@ def wds_write_worker(
     count = 0
     buf = {}
     more = True
-    with tqdm(total=total, initial=0, unit="images", unit_scale=True, leave=False) as progress:
-        while more:
-            deq: Optional[tuple[int, bytes, str, int]] = q_out.get()
-            if deq is not None:
-                (idx, sample, suffix, target) = deq
-                buf[idx] = (sample, suffix, target)
+    try:
+        with tqdm(total=total, initial=0, unit="images", unit_scale=True, leave=False) as progress:
+            while more:
+                deq: Optional[tuple[int, bytes, str, int]] = q_out.get()
+                if deq is not None:
+                    (idx, sample, suffix, target) = deq
+                    buf[idx] = (sample, suffix, target)
 
-            else:
-                more = False
+                else:
+                    more = False
 
-            # Ensures ordered write
-            while count in buf:
-                (sample, suffix, target) = buf[count]
-                del buf[count]
+                # Ensures ordered write
+                while count in buf:
+                    (sample, suffix, target) = buf[count]
+                    del buf[count]
 
-                if args.no_cls is True:
-                    cls = {}
-                else:
-                    cls = {"cls": target}
+                    if args.no_cls is True:
+                        cls = {}
+                    else:
+                        cls = {"cls": target}
 
-                sink.write(
-                    {
-                        "__key__": f"sample{count:09d}",
-                        suffix: sample,
-                        **cls,
-                    }
-                )
+                    sink.write(
+                        {
+                            "__key__": f"sample{count:09d}",
+                            suffix: sample,
+                            **cls,
+                        }
+                    )
 
-                count += 1
+                    count += 1
 
-                # Update progress bar
-                progress.update(n=1)
+                    # Update progress bar
+                    progress.update(n=1)
+
+    except Exception:
+        error_event.set()
+        raise
 
     sink.close()
 
@@ -218,35 +228,42 @@ def wds_write_worker(
 
 
 def directory_write_worker(
-    q_out: Any, _error_event: Any, pack_path: Path, total: int, _: argparse.Namespace, idx_to_class: dict[int, str]
+    q_out: Any, error_event: Any, pack_path: Path, total: int, _: argparse.Namespace, idx_to_class: dict[int, str]
 ) -> None:
     count = 0
     buf = {}
     more = True
-    with tqdm(total=total, initial=0, unit="images", unit_scale=True, leave=False) as progress:
-        while more:
-            deq: Optional[tuple[int, bytes, str, int]] = q_out.get()
-            if deq is not None:
-                (idx, sample, suffix, target) = deq
-                buf[idx] = (sample, suffix, target)
+    try:
+        with tqdm(total=total, initial=0, unit="images", unit_scale=True, leave=False) as progress:
+            while more:
+                deq: Optional[tuple[int, bytes, str, int]] = q_out.get()
+                if deq is not None:
+                    (idx, sample, suffix, target) = deq
+                    buf[idx] = (sample, suffix, target)
 
-            else:
-                more = False
+                else:
+                    more = False
+
+                # Ensures ordered write
+                while count in buf:
+                    (sample, suffix, target) = buf[count]
+                    del buf[count]
+                    with open(
+                        pack_path.joinpath(idx_to_class[target]).joinpath(f"{count:06d}.{suffix}"), "wb"
+                    ) as handle:
+                        handle.write(sample)
 
-            # Ensures ordered write
-            while count in buf:
-                (sample, suffix, target) = buf[count]
-                del buf[count]
-                with open(pack_path.joinpath(idx_to_class[target]).joinpath(f"{count:06d}.{suffix}"), "wb") as handle:
-                    handle.write(sample)
+                    count += 1
 
-                count += 1
+                    # Update progress bar
+                    progress.update(n=1)
 
-                # Update progress bar
-                progress.update(n=1)
+    except Exception:
+        error_event.set()
+        raise
 
 
-# pylint: disable=too-many-locals,too-many-branches
+# pylint: disable=too-many-locals,too-many-branches,too-many-statements
 def pack(args: argparse.Namespace, pack_path: Path) -> None:
     if args.sampling_file is not None:
         with open(args.sampling_file, "r", encoding="utf-8") as handle:
@@ -308,9 +325,7 @@ def pack(args: argparse.Namespace, pack_path: Path) -> None:
     read_processes: list[multiprocessing.Process] = []
     for idx in range(args.jobs):
         read_processes.append(
-            multiprocessing.Process(
-                target=read_worker, args=(q_in[idx], q_out, error_event, args.size, args.format), daemon=True
-            )
+            multiprocessing.Process(target=read_worker, args=(q_in[idx], q_out, error_event, args.size, args.format))
         )
 
     for p in read_processes:
@@ -326,53 +341,107 @@ def pack(args: argparse.Namespace, pack_path: Path) -> None:
         raise ValueError("Unknown pack type")
 
     write_process = multiprocessing.Process(
-        target=target_writer, args=(q_out, error_event, pack_path, len(dataset), args, idx_to_class), daemon=True
+        target=target_writer, args=(q_out, error_event, pack_path, len(dataset), args, idx_to_class)
     )
     write_process.start()
 
+    # Flag to prevent signal handler re-entry
+    cleanup_in_progress = False
+
+    def cleanup_processes() -> None:
+        nonlocal cleanup_in_progress
+        if cleanup_in_progress is True:
+            return
+
+        cleanup_in_progress = True
+
+        # Cancel queue join threads to prevent blocking during cleanup
+        for q in q_in:
+            q.cancel_join_thread()
+
+        q_out.cancel_join_thread()
+
+        # Terminate child processes
+        for p in read_processes:
+            if p.is_alive():
+                p.terminate()
+
+        if write_process.is_alive():
+            write_process.terminate()
+
+        # Wait briefly for termination
+        for p in read_processes:
+            p.join(timeout=1)
+
+        write_process.join(timeout=1)
+
     def signal_handler(signum, _frame) -> None:  # type: ignore
         logger.info(f"Received signal: {signum} at {multiprocessing.current_process().name}, aborting...")
         error_event.set()
+        cleanup_processes()
         raise SystemExit(1)
 
     signal.signal(signal.SIGINT, signal_handler)
 
-    tic = time.time()
-    for idx, sample_idx in enumerate(indices):
-        if idx % 1000 == 0:
-            if error_event.is_set() is True:
-                raise RuntimeError()
-
-        (path, target) = dataset[sample_idx]
-        q_in[idx % len(q_in)].put((idx, path, target), block=True, timeout=None)
-
-    for q in q_in:
-        q.put(None, block=True, timeout=None)
-
-    for p in read_processes:
+    try:
+        tic = time.time()
+        for idx, sample_idx in enumerate(indices):
+            if idx % 1000 == 0:
+                if error_event.is_set() is True:
+                    cleanup_processes()
+                    raise RuntimeError()
+
+            (path, target) = dataset[sample_idx]
+
+            while True:
+                try:
+                    q_in[idx % len(q_in)].put((idx, path, target), block=True, timeout=1)
+                    break
+                except queue.Full:
+                    if error_event.is_set() is True:
+                        cleanup_processes()
+                        raise RuntimeError()  # pylint: disable=raise-missing-from
+
+        for q in q_in:
+            q.put(None, block=True, timeout=None)
+
+        for p in read_processes:
+            while True:
+                p.join(timeout=2)
+                if p.is_alive() is False:
+                    break
+
+                if error_event.is_set() is True:
+                    cleanup_processes()
+                    raise RuntimeError()
+
+        q_out.put(None, block=True, timeout=None)
         while True:
-            p.join(timeout=2)
-            if p.is_alive() is False:
+            write_process.join(timeout=2)
+            if write_process.is_alive() is False:
                 break
 
             if error_event.is_set() is True:
+                cleanup_processes()
                 raise RuntimeError()
 
-    q_out.put(None, block=True, timeout=None)
-    write_process.join()
+        if error_event.is_set() is True:
+            cleanup_processes()
+            raise RuntimeError()
 
-    if error_event.is_set() is True:
-        raise RuntimeError()
+        if args.type == "wds":
+            (wds_path, num_shards) = fs_ops.wds_braces_from_path(pack_path, prefix=f"{args.suffix}-{args.split}")
+            logger.info(f"Packed {len(dataset):,} samples into {num_shards} shards at {wds_path}")
+        elif args.type == "directory":
+            logger.info(f"Packed {len(dataset):,} samples")
 
-    if args.type == "wds":
-        (wds_path, num_shards) = fs_ops.wds_braces_from_path(pack_path, prefix=f"{args.suffix}-{args.split}")
-        logger.info(f"Packed {len(dataset):,} samples into {num_shards} shards at {wds_path}")
-    elif args.type == "directory":
-        logger.info(f"Packed {len(dataset):,} samples")
+        toc = time.time()
+        rate = len(dataset) / (toc - tic)
+        logger.info(f"{format_duration(toc-tic)} to pack {len(dataset):,} samples ({rate:.2f} samples/sec)")
 
-    toc = time.time()
-    rate = len(dataset) / (toc - tic)
-    logger.info(f"{format_duration(toc-tic)} to pack {len(dataset):,} samples ({rate:.2f} samples/sec)")
+    except Exception:
+        cleanup_processes()
+        raise
 
 
 def set_parser(subparsers: Any) -> None:

birder/tools/show_det_iterator.py

@@ -43,6 +43,7 @@ def show_det_iterator(args: argparse.Namespace) -> None:
             args.dynamic_size,
             args.multiscale,
             args.max_size,
+            args.multiscale_min_size,
         )
         mosaic_transforms = training_preset(
             args.size,
@@ -52,6 +53,7 @@ def show_det_iterator(args: argparse.Namespace) -> None:
             args.dynamic_size,
             args.multiscale,
             args.max_size,
+            args.multiscale_min_size,
             post_mosaic=True,
         )
         if args.mosaic_prob > 0.0:
@@ -160,7 +162,9 @@ def show_det_iterator(args: argparse.Namespace) -> None:
 
     else:
         if args.batch_multiscale is True:
-            data_collate_fn: Any = BatchRandomResizeCollator(offset, args.size)
+            data_collate_fn: Any = BatchRandomResizeCollator(
+                offset, args.size, multiscale_min_size=args.multiscale_min_size
+            )
         else:
             data_collate_fn = collate_fn
 
@@ -259,6 +263,11 @@ def set_parser(subparsers: Any) -> None:
         help="allow variable image sizes while preserving aspect ratios",
     )
     subparser.add_argument("--multiscale", default=False, action="store_true", help="enable random scale per image")
+    subparser.add_argument(
+        "--multiscale-min-size",
+        type=int,
+        help="minimum short-edge size for multiscale lists (rounded up to nearest multiple of 32)",
+    )
     subparser.add_argument(
         "--batch-multiscale",
         default=False,

birder/version.py

@@ -1 +1 @@
-__version__ = "v0.2.2"
+__version__ = "v0.2.3"