birder-clip 0.0.2.dev4__tar.gz → 0.0.2.dev5__tar.gz

{birder_clip-0.0.2.dev4 → birder_clip-0.0.2.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: birder_clip
-Version: 0.0.2.dev4
+Version: 0.0.2.dev5
 Summary: A Birder extension for CLIP-style image-text modeling and multimodal computer vision workflows.
 Author: Ofer Hasson
 License-Expression: Apache-2.0
@@ -24,7 +24,7 @@ Classifier: Typing :: Typed
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: birder>=0.5.4
+Requires-Dist: birder>=0.5.6
 Requires-Dist: ftfy>=6.3.1
 Requires-Dist: regex>=2025.7.29
 Requires-Dist: tqdm>=4.67.0
@@ -37,7 +37,7 @@ Provides-Extra: dev
 Requires-Dist: bandit~=1.9.4; extra == "dev"
 Requires-Dist: black~=26.5.0; extra == "dev"
 Requires-Dist: build~=1.5.0; extra == "dev"
-Requires-Dist: bumpver~=2025.1131; extra == "dev"
+Requires-Dist: bumpver~=2026.1132; extra == "dev"
 Requires-Dist: coverage~=7.14.1; extra == "dev"
 Requires-Dist: debugpy; extra == "dev"
 Requires-Dist: flake8-pep585~=0.1.7; extra == "dev"

{birder_clip-0.0.2.dev4 → birder_clip-0.0.2.dev5}/birder_clip/common/fs_ops.py

@@ -67,7 +67,7 @@ def model_path(
     network_name: str,
     *,
     epoch: Optional[int | str] = None,
-    file_format: FileFormatType = "pt",
+    st: bool = False,
     states: bool = False,
 ) -> Path:
     if epoch is not None:
@@ -77,8 +77,10 @@ def model_path(
 
     if states is True:
         file_name = f"{file_name}_states.pt"
+    elif st is True:
+        file_name = f"{file_name}.safetensors"
     else:
-        file_name = f"{file_name}.{file_format}"
+        file_name = f"{file_name}.pt"
 
     return settings.MODELS_DIR.joinpath(file_name)
 
@@ -109,6 +111,30 @@ def _checkpoint_states(
     torch.save(kwargs, states_path)
 
 
+def _checkpoint_states_from_state_dicts(
+    states_path: Path,
+    optimizer_state: Optional[dict[str, Any]],
+    scheduler_state: Optional[dict[str, Any]],
+    scaler_state: Optional[dict[str, Any]],
+    model_base_state: Optional[dict[str, Any]],
+    **extra_states: Optional[dict[str, Any]],
+) -> None:
+    if optimizer_state is None or scheduler_state is None:
+        return
+
+    logger.info(f"Saving checkpoint states {states_path}...")
+    torch.save(
+        {
+            "optimizer_state": optimizer_state,
+            "scheduler_state": scheduler_state,
+            "scaler_state": scaler_state,
+            "model_base_state": model_base_state,
+            **extra_states,
+        },
+        states_path,
+    )
+
+
 class TrainingStates(NamedTuple):
     optimizer_state: Optional[dict[str, Any]]
     scheduler_state: Optional[dict[str, Any]]
@@ -182,6 +208,50 @@ def checkpoint_model(
     _checkpoint_states(states_path, optimizer, scheduler, scaler, model_base, **extra_states)
 
 
+def checkpoint_model_from_state_dicts(
+    network_name: str,
+    epoch: int,
+    model_state: dict[str, Any],
+    task: Any,
+    signature: SignatureType,
+    rgb_stats: RGBType,
+    optimizer_state: Optional[dict[str, Any]],
+    scheduler_state: Optional[dict[str, Any]],
+    scaler_state: Optional[dict[str, Any]],
+    model_base_state: Optional[dict[str, Any]],
+    *,
+    external_config: Optional[dict[str, Any]] = None,
+    **extra_states: Optional[dict[str, Any]],
+) -> None:
+    kwargs = {}
+    if external_config is not None:
+        kwargs["config"] = external_config
+
+    path = model_path(network_name, epoch=epoch)
+    states_path = model_path(network_name, epoch=epoch, states=True)
+    logger.info(f"Saving model checkpoint {path}...")
+    torch.save(
+        {
+            "state": model_state,
+            "birder_clip_version": __version__,
+            "task": task,
+            "signature": signature,
+            "rgb_stats": rgb_stats,
+            **kwargs,
+        },
+        path,
+    )
+
+    _checkpoint_states_from_state_dicts(
+        states_path,
+        optimizer_state,
+        scheduler_state,
+        scaler_state,
+        model_base_state,
+        **extra_states,
+    )
+
+
 def clean_checkpoints(network_name: str, keep_last: int) -> None:
     epoch = "*[0-9]"
     models_glob = str(model_path(network_name, epoch=epoch))
@@ -314,7 +384,7 @@ def load_model(
             embed_dim=embed_dim,
             tokenizer=tokenizer,
         )
-        path = model_path(_network_name, epoch=epoch, file_format="safetensors" if st is True else "pt")
+        path = model_path(_network_name, epoch=epoch, st=st)
 
     logger.info(f"Loading model from {path} on device {device}...")
 
@@ -589,6 +659,33 @@ def load_pretrained_tokenizer(weights: str, *, download: bool = True, **kwargs:
     return get_tokenizer(tokenizer_name, **tokenizer_kwargs)
 
 
+def save_st(
+    net: torch.nn.Module,
+    dst: str,
+    task: str,
+    signature: SignatureType,
+    rgb_stats: RGBType,
+    *,
+    external_config: Optional[dict[str, Any]] = None,
+) -> None:
+    assert _HAS_SAFETENSORS, "'pip install safetensors' to use .safetensors"
+    kwargs = {}
+    if external_config is not None:
+        kwargs["config"] = json.dumps(external_config)
+
+    safetensors.torch.save_model(
+        net,
+        str(dst),
+        {
+            "birder_clip_version": __version__,
+            "task": task,
+            "signature": json.dumps(signature),
+            "rgb_stats": json.dumps(rgb_stats),
+            **kwargs,
+        },
+    )
+
+
 def download_model_by_weights(
     weights: str, *, dst: Optional[str | Path] = None, file_format: FileFormatType = "pt", progress_bar: bool = True
 ) -> None:

{birder_clip-0.0.2.dev4 → birder_clip-0.0.2.dev5}/birder_clip/common/lib.py

@@ -86,7 +86,7 @@ def get_image_text_model_config(
 
     if config is not None:
         for key, value in config.items():
-            if key in {"image", "text"} and isinstance(value, dict) is True:
+            if key in {"image", "text"} and isinstance(value, dict):
                 model_config[key] = {**model_config.get(key, {}), **value}
             else:
                 model_config[key] = value

{birder_clip-0.0.2.dev4 → birder_clip-0.0.2.dev5}/birder_clip/common/training_cli.py

@@ -49,7 +49,9 @@ def add_loss_args(parser: argparse.ArgumentParser) -> None:
 def add_optimization_args(parser: argparse.ArgumentParser, default_batch_size: int = 32) -> None:
     group = parser.add_argument_group("Optimization parameters")
     group.add_argument("--batch-size", type=int, default=default_batch_size, metavar="N", help="the batch size")
-    group.add_argument("--opt", type=str, choices=list(get_args(OptimizerType)), default="sgd", help="optimizer to use")
+    group.add_argument(
+        "--opt", type=str, choices=list(get_args(OptimizerType)), default="adamw", help="optimizer to use"
+    )
     group.add_argument("--opt-fused", default=False, action="store_true", help="use fused optimizer implementation")
     group.add_argument("--momentum", type=float, default=0.9, metavar="M", help="optimizer momentum")
     group.add_argument("--nesterov", default=False, action="store_true", help="use nesterov momentum")
@@ -318,7 +320,13 @@ def add_dataloader_args(parser: argparse.ArgumentParser) -> None:
         action="store_true",
         help="keep dataloader worker processes alive between epochs",
     )
-    group.add_argument("--drop-last", default=False, action="store_true", help="drop the last incomplete batch")
+    group.add_argument(
+        "--no-drop-last",
+        dest="drop_last",
+        default=True,
+        action="store_false",
+        help="do not drop the last incomplete batch",
+    )
 
 
 def add_precision_args(parser: argparse.ArgumentParser) -> None:
@@ -410,6 +418,44 @@ def add_distributed_args(parser: argparse.ArgumentParser) -> None:
     group.add_argument("--local-rank", type=int, metavar="N", help="local rank")
     group.add_argument("--dist-url", type=str, default="env://", help="URL used to initialize distributed training")
     group.add_argument("--dist-backend", type=str, default="nccl", help="distributed backend")
+    group.add_argument(
+        "--distributed-mode", type=str, choices=["ddp", "fsdp"], default="ddp", help="distributed training mode"
+    )
+    group.add_argument(
+        "--fsdp-sharding-strategy",
+        type=str,
+        choices=["shard-grad-op", "full-shard"],
+        default="shard-grad-op",
+        help="FSDP sharding strategy",
+    )
+    group.add_argument(
+        "--fsdp-param-dtype",
+        type=str,
+        choices=["float32", "float16", "bfloat16"],
+        help="FSDP mixed precision parameter dtype",
+    )
+    group.add_argument(
+        "--fsdp-reduce-dtype",
+        type=str,
+        choices=["float32", "float16", "bfloat16"],
+        help="FSDP mixed precision gradient reduction dtype",
+    )
+    group.add_argument(
+        "--fsdp-wrap-policy",
+        type=str,
+        choices=["block-group-regex", "min-num-params"],
+        default="block-group-regex",
+        help="FSDP module wrapping policy",
+    )
+    group.add_argument(
+        "--fsdp-wrap-min-num-params",
+        type=float,
+        metavar="M",
+        help="minimum module parameter count in millions for wrapping when using --fsdp-wrap-policy min-num-params",
+    )
+    group.add_argument(
+        "--fsdp-offload-policy", type=str, choices=["none", "cpu"], default="none", help="FSDP parameter offload policy"
+    )
     group.add_argument(
         "--find-unused-parameters",
         default=False,
@@ -561,3 +607,23 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--grad-accum-steps must be >= 1")
     if args.model_ema_steps < 1:
         raise cli.ValidationError("--model-ema-steps must be >= 1")
+
+    if args.distributed_mode == "fsdp":
+        if args.sync_bn is True:
+            raise cli.ValidationError("--sync-bn cannot be used with --distributed-mode fsdp")
+        if args.find_unused_parameters is True:
+            raise cli.ValidationError("--find-unused-parameters cannot be used with --distributed-mode fsdp")
+        if args.compile_opt is True:
+            raise cli.ValidationError("--compile-opt cannot be used with --distributed-mode fsdp")
+        if args.compile_fullgraph is True:
+            raise cli.ValidationError("--compile-fullgraph cannot be used with --distributed-mode fsdp")
+        if args.cpu is True:
+            raise cli.ValidationError("--cpu cannot be used with --distributed-mode fsdp")
+        if args.model_ema is True:
+            raise cli.ValidationError("--model-ema cannot be used with --distributed-mode fsdp")
+        if args.fsdp_wrap_policy == "min-num-params" and args.fsdp_wrap_min_num_params is None:
+            raise cli.ValidationError(
+                "--fsdp-wrap-min-num-params is required when --fsdp-wrap-policy is min-num-params"
+            )
+        if args.fsdp_wrap_min_num_params is not None and args.fsdp_wrap_min_num_params <= 0:
+            raise cli.ValidationError("--fsdp-wrap-min-num-params must be > 0")

birder_clip-0.0.2.dev5/birder_clip/common/training_utils.py

@@ -0,0 +1,99 @@
+import argparse
+import logging
+from pathlib import Path
+from typing import Any
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from birder.common import fsdp_utils
+from birder.common import training_utils as birder_training_utils
+
+from birder_clip.common import fs_ops
+from birder_clip.conf import settings
+
+
+def setup_file_logging(log_file_path: str | Path) -> logging.Handler:
+    file_handler = logging.FileHandler(log_file_path)
+    formatter = logging.Formatter(
+        fmt="{message}",
+        style="{",
+    )
+    file_handler.setFormatter(formatter)
+    file_handler.setLevel(settings.LOG_LEVEL)
+
+    logging.getLogger("birder").addHandler(file_handler)
+    logging.getLogger("birder_clip").addHandler(file_handler)
+
+    return file_handler
+
+
+def save_training_checkpoint(
+    args: argparse.Namespace,
+    network_name: str,
+    epoch: int,
+    net: torch.nn.Module,
+    signature: Any,
+    rgb_stats: Any,
+    optimizer: Optional[torch.optim.Optimizer],
+    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
+    scaler: Optional[torch.amp.grad_scaler.GradScaler],
+    model_base: Optional[torch.nn.Module],
+    *,
+    fsdp_mode: bool = False,
+    fsdp_model_state: Optional[dict[str, Any]] = None,
+    external_config: Optional[dict[str, Any]] = None,
+    **extra_states: Optional[dict[str, Any]],
+) -> None:
+    if fsdp_mode is True:
+        if fsdp_model_state is not None:
+            model_state = fsdp_model_state
+        else:
+            model_state = fsdp_utils.gather_full_model_state_dict(net)
+
+        optimizer_state = None
+        scheduler_state = None
+        scaler_state = None
+        model_base_state = None
+        if optimizer is not None and scheduler is not None:
+            optimizer_state = fsdp_utils.gather_full_optimizer_state_dict(net, optimizer)
+            scheduler_state = scheduler.state_dict()
+            if scaler is not None:
+                scaler_state = scaler.state_dict()
+            if model_base is not None:
+                model_base_state = model_base.state_dict()
+
+        if birder_training_utils.is_global_primary(args) is True:
+            fs_ops.checkpoint_model_from_state_dicts(
+                network_name,
+                epoch,
+                model_state,
+                net.task,
+                signature,
+                rgb_stats,
+                optimizer_state,
+                scheduler_state,
+                scaler_state,
+                model_base_state,
+                external_config=external_config,
+                **extra_states,
+            )
+
+        if birder_training_utils.is_dist_available_and_initialized() is True:
+            dist.barrier()
+
+    else:
+        if birder_training_utils.is_global_primary(args) is True:
+            fs_ops.checkpoint_model(
+                network_name,
+                epoch,
+                net,
+                signature,
+                rgb_stats,
+                optimizer,
+                scheduler,
+                scaler,
+                model_base,
+                external_config=external_config,
+                **extra_states,
+            )

{birder_clip-0.0.2.dev4 → birder_clip-0.0.2.dev5}/birder_clip/data/datasets/webdataset.py

@@ -24,10 +24,10 @@ def decode_caption(caption: Any, caption_json_key: str = "caption") -> str:
     if isinstance(caption, bytes):
         caption = caption.decode("utf-8")
 
-    if isinstance(caption, str) is False:
+    if not isinstance(caption, str):
         raise TypeError(f"WebDataset caption must be a string, got {type(caption).__name__}")
 
-    return caption  # type: ignore[no-any-return]
+    return caption
 
 
 def tokenize_caption(caption: str, tokenizer: Tokenizer) -> torch.Tensor:

birder_clip-0.0.2.dev5/birder_clip/inference/zero_shot.py

@@ -0,0 +1,128 @@
+"""
+Zero-shot text embedding helpers
+
+Zero-shot classification compares image features against one text feature per
+candidate class. When multiple prompt templates are used, this module follows
+the OpenCLIP/OpenAI CLIP convention: encode every class/template prompt,
+normalize prompt embeddings, average them per class and normalize the averaged
+class embedding again.
+"""
+
+import sys
+from collections.abc import Callable
+from collections.abc import Iterator
+from collections.abc import Sequence
+from typing import Optional
+
+import numpy as np
+import numpy.typing as npt
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from birder_clip.net.base import BaseNet
+from birder_clip.tokenizers.base import Tokenizer
+
+
+def render_prompts(class_names: Sequence[str], templates: Sequence[str]) -> list[str]:
+    return [template.format(class_name) for class_name in class_names for template in templates]
+
+
+def build_class_text_embeddings(
+    net: BaseNet,
+    tokenizer: Tokenizer,
+    class_names: Sequence[str],
+    templates: Sequence[str],
+    *,
+    device: torch.device,
+    context_length: Optional[int] = None,
+    batch_size: Optional[int] = None,
+    amp: bool = False,
+    amp_dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    num_templates = len(templates)
+    if batch_size is None:
+        batch_size = len(class_names)
+
+    class_text_embeddings = []
+    with torch.inference_mode():
+        for start in range(0, len(class_names), batch_size):
+            batch_class_names = class_names[start : start + batch_size]
+            prompts = render_prompts(batch_class_names, templates)
+            tokens = tokenizer(prompts, context_length=context_length).to(device)
+            with torch.amp.autocast(device.type, enabled=amp, dtype=amp_dtype):
+                class_embeddings = net.encode_text(tokens, normalize=True)
+
+            class_embeddings = class_embeddings.reshape(len(batch_class_names), num_templates, -1).mean(dim=1)
+            class_embeddings = F.normalize(class_embeddings, dim=-1)
+            class_text_embeddings.append(class_embeddings)
+
+    return torch.concat(class_text_embeddings, dim=0)
+
+
+DataloaderInferenceResult = tuple[list[str], npt.NDArray[np.float32], npt.NDArray[np.int64]]
+
+
+def infer_dataloader_iter(
+    device: torch.device,
+    net: BaseNet | torch.ScriptModule,
+    dataloader: DataLoader,
+    text_embeddings: torch.Tensor,
+    return_logits: bool = False,
+    model_dtype: torch.dtype = torch.float32,
+    amp: bool = False,
+    amp_dtype: Optional[torch.dtype] = None,
+    num_samples: Optional[int] = None,
+    batch_callback: Optional[Callable[[list[str], npt.NDArray[np.float32], npt.NDArray[np.int64]], None]] = None,
+    chunk_size: Optional[float] = None,
+) -> Iterator[DataloaderInferenceResult]:
+    if chunk_size is None:
+        chunk_size = float("inf")
+
+    net.to(device, dtype=model_dtype)
+    out_list: list[npt.NDArray[np.float32]] = []
+    labels_list: list[npt.NDArray[np.int64]] = []
+    sample_paths: list[str] = []
+    sample_count = 0
+    with tqdm(total=num_samples, initial=0, unit="images", unit_scale=True, leave=False) as progress:
+        for file_paths, inputs, targets in dataloader:
+            batch_size = inputs.size(0)
+
+            # Inference
+            inputs = inputs.to(device, dtype=model_dtype)
+            with torch.amp.autocast(device.type, enabled=amp, dtype=amp_dtype):
+                image_embeddings = net.encode_image(inputs, normalize=True)
+                logits = net.forward_logits(image_embeddings, text_embeddings)
+                if return_logits is True:
+                    out = logits.cpu().float().numpy()
+                else:
+                    out = F.softmax(logits, dim=-1).cpu().float().numpy()
+
+            out_list.append(out)
+
+            # Set labels and sample list
+            batch_labels = targets.cpu().numpy()
+            labels_list.append(batch_labels)
+            sample_paths.extend(file_paths)
+
+            if batch_callback is not None:
+                batch_callback(file_paths, out, batch_labels)
+
+            # Update progress bar
+            progress.update(n=batch_size)
+
+            # Yield results when we reach chunk_size
+            sample_count += batch_size
+            if sample_count >= chunk_size:
+                with tqdm.external_write_mode(file=sys.stderr):
+                    yield (sample_paths, np.concatenate(out_list, axis=0), np.concatenate(labels_list))
+
+                # Reset for next chunk
+                out_list = []
+                labels_list = []
+                sample_paths = []
+                sample_count = 0
+
+    if len(out_list) > 0:
+        yield (sample_paths, np.concatenate(out_list, axis=0), np.concatenate(labels_list))