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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: birder_clip
-Version: 0.0.2.dev3
+Version: 0.0.2.dev4
 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.2
+Requires-Dist: birder>=0.5.4
 Requires-Dist: ftfy>=6.3.1
 Requires-Dist: regex>=2025.7.29
 Requires-Dist: tqdm>=4.67.0

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

@@ -43,9 +43,17 @@ def get_image_text_network_name(
     parts = [network]
     if image_encoder is not None:
         parts.append(image_encoder)
-    if text_encoder is not None:
+    if text_encoder is not None and text_encoder != "text_transformer":
         parts.append(text_encoder)
-    if tokenizer is not None:
+
+    if registry.exists(network) is True:
+        default_tokenizer = registry.get_default_tokenizer(network)
+    else:
+        default_tokenizer = "simple_tokenizer"
+    if default_tokenizer is None:
+        default_tokenizer = "simple_tokenizer"
+
+    if tokenizer is not None and tokenizer != default_tokenizer:
         parts.append(tokenizer)
     if embed_dim is not None:
         parts.append(f"d{embed_dim}")

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

@@ -17,32 +17,6 @@ from birder_clip.model_registry import Task
 from birder_clip.model_registry import registry
 
 
-def add_compile_args(parser: argparse.ArgumentParser) -> None:
-    group = parser.add_argument_group("Compilation parameters")
-    group.add_argument("--compile", default=False, action="store_true", help="enable compilation")
-    group.add_argument("--compile-fullgraph", default=False, action="store_true", help="compile using fullgraph=True")
-    group.add_argument(
-        "--compile-mode", type=str, choices=list(torch._inductor.list_mode_options().keys()), help="torch.compile mode"
-    )
-    group.add_argument(
-        "--compile-opt", default=False, action="store_true", help="enable compilation for optimizer step"
-    )
-    group.add_argument(
-        "--compile-recompile-limit",
-        type=int,
-        default=torch.compiler.config.recompile_limit,
-        metavar="N",
-        help="maximum recompilations per compiled function before eager fallback",
-    )
-    group.add_argument(
-        "--compile-accumulated-recompile-limit",
-        type=int,
-        default=torch.compiler.config.accumulated_recompile_limit,
-        metavar="N",
-        help="maximum total recompilations across compiled functions",
-    )
-
-
 def add_model_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("-n", "--network", type=str, help="the image-text network to train")
     parser.add_argument("-t", "--tag", type=str, help="add model tag")
@@ -75,9 +49,7 @@ 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="adamw", help="optimizer to use"
-    )
+    group.add_argument("--opt", type=str, choices=list(get_args(OptimizerType)), default="sgd", 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")
@@ -92,6 +64,7 @@ def add_optimization_args(parser: argparse.ArgumentParser, default_batch_size: i
         metavar="N",
         help="number of iterations to accumulate gradients per optimizer step",
     )
+    # NOTE: Add flag for negative sample caching in grad accum mode
 
 
 def add_lr_wd_args(parser: argparse.ArgumentParser) -> None:
@@ -129,14 +102,14 @@ def add_lr_scheduler_args(parser: argparse.ArgumentParser) -> None:
         "--lr-scheduler-update",
         type=str,
         choices=["epoch", "step"],
-        default="step",
+        default="epoch",
         help="when to apply learning rate scheduler update: epoch (once per epoch), step (each optimizer step)",
     )
     group.add_argument(
         "--lr-scheduler",
         type=str,
         choices=list(get_args(SchedulerType)),
-        default="cosine",
+        default="constant",
         help="learning rate scheduler",
     )
     group.add_argument(
@@ -175,15 +148,6 @@ def add_lr_scheduler_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
-def add_input_args(parser: argparse.ArgumentParser) -> None:
-    group = parser.add_argument_group("Input parameters")
-    group.add_argument(
-        "--channels", type=int, default=settings.DEFAULT_NUM_CHANNELS, metavar="N", help="no. of image channels"
-    )
-    group.add_argument("--size", type=int, nargs="+", metavar=("H", "W"), help="image size")
-    group.add_argument("--context-length", type=int, metavar="N", help="text context length")
-
-
 def add_training_schedule_args(parser: argparse.ArgumentParser, default_epochs: int = 100) -> None:
     group = parser.add_argument_group("Training schedule parameters")
     group.add_argument("--epochs", type=int, default=default_epochs, metavar="N", help="number of training epochs")
@@ -204,6 +168,37 @@ def add_training_schedule_args(parser: argparse.ArgumentParser, default_epochs:
     )
 
 
+def add_ema_args(
+    parser: argparse.ArgumentParser, default_ema_steps: int = 1, default_ema_decay: float = 0.9999
+) -> None:
+    group = parser.add_argument_group("Exponential moving average parameters")
+    group.add_argument(
+        "--model-ema",
+        default=False,
+        action="store_true",
+        help="enable tracking exponential moving average of model parameters",
+    )
+    group.add_argument(
+        "--model-ema-steps",
+        type=int,
+        default=default_ema_steps,
+        metavar="N",
+        help="number of optimizer steps between EMA updates",
+    )
+    group.add_argument(
+        "--model-ema-decay",
+        type=float,
+        default=default_ema_decay,
+        help="decay factor for exponential moving average of model parameters",
+    )
+    group.add_argument(
+        "--model-ema-warmup",
+        type=int,
+        metavar="N",
+        help="number of epochs/steps before EMA is applied (defaults to warmup epochs/steps, pass 0 to disable warmup)",
+    )
+
+
 def add_batch_norm_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Batch normalization parameters")
     group.add_argument(
@@ -215,6 +210,15 @@ def add_batch_norm_args(parser: argparse.ArgumentParser) -> None:
     group.add_argument("--sync-bn", default=False, action="store_true", help="use synchronized BatchNorm")
 
 
+def add_input_args(parser: argparse.ArgumentParser) -> None:
+    group = parser.add_argument_group("Input parameters")
+    group.add_argument(
+        "--channels", type=int, default=settings.DEFAULT_NUM_CHANNELS, metavar="N", help="no. of image channels"
+    )
+    group.add_argument("--size", type=int, nargs="+", metavar=("H", "W"), help="image size")
+    group.add_argument("--context-length", type=int, metavar="N", help="text context length")
+
+
 def add_data_aug_args(
     parser: argparse.ArgumentParser,
     default_level: int = 4,
@@ -260,7 +264,7 @@ def add_data_aug_args(
         "--rgb-mode",
         type=str,
         choices=list(typing.get_args(RGBMode)),
-        default="clip",
+        default="birder",
         help="RGB mean and std to use for normalization",
     )
     group.add_argument(
@@ -279,67 +283,6 @@ def add_data_aug_args(
     )
 
 
-def add_checkpoint_args(parser: argparse.ArgumentParser, default_save_frequency: int = 1) -> None:
-    group = parser.add_argument_group("Checkpoint parameters")
-    group.add_argument(
-        "--save-frequency", type=int, default=default_save_frequency, metavar="N", help="frequency of model saving"
-    )
-    group.add_argument(
-        "--keep-last", type=int, metavar="N", help="number of recent checkpoints to keep (older ones are deleted)"
-    )
-    group.add_argument(
-        "--pretrained",
-        default=False,
-        action="store_true",
-        help="start with pretrained version of specified network (will download if not found locally)",
-    )
-    group.add_argument("--resume-epoch", type=int, metavar="N", help="epoch number to resume training from")
-    group.add_argument(
-        "--non-strict-weights",
-        default=False,
-        action="store_true",
-        help="allow non-strict loading of model weights (missing or unexpected keys in state_dict)",
-    )
-    group.add_argument(
-        "--load-states",
-        default=False,
-        action="store_true",
-        help="load optimizer, scheduler and scaler states when resuming",
-    )
-    group.add_argument("--load-scheduler", default=False, action="store_true", help="load only scheduler when resuming")
-
-
-def add_ema_args(
-    parser: argparse.ArgumentParser, default_ema_steps: int = 1, default_ema_decay: float = 0.9999
-) -> None:
-    group = parser.add_argument_group("Exponential moving average parameters")
-    group.add_argument(
-        "--model-ema",
-        default=False,
-        action="store_true",
-        help="enable tracking exponential moving average of model parameters",
-    )
-    group.add_argument(
-        "--model-ema-steps",
-        type=int,
-        default=default_ema_steps,
-        metavar="N",
-        help="number of optimizer steps between EMA updates",
-    )
-    group.add_argument(
-        "--model-ema-decay",
-        type=float,
-        default=default_ema_decay,
-        help="decay factor for exponential moving average of model parameters",
-    )
-    group.add_argument(
-        "--model-ema-warmup",
-        type=int,
-        metavar="N",
-        help="number of epochs/steps before EMA is applied (defaults to warmup epochs/steps, pass 0 to disable warmup)",
-    )
-
-
 def add_dataloader_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Dataloader parameters")
     group.add_argument(
@@ -405,6 +348,62 @@ def add_precision_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
+def add_compile_args(parser: argparse.ArgumentParser) -> None:
+    group = parser.add_argument_group("Compilation parameters")
+    group.add_argument("--compile", default=False, action="store_true", help="enable compilation")
+    group.add_argument("--compile-fullgraph", default=False, action="store_true", help="compile using fullgraph=True")
+    group.add_argument(
+        "--compile-mode", type=str, choices=list(torch._inductor.list_mode_options().keys()), help="torch.compile mode"
+    )
+    group.add_argument(
+        "--compile-opt", default=False, action="store_true", help="enable compilation for optimizer step"
+    )
+    group.add_argument(
+        "--compile-recompile-limit",
+        type=int,
+        default=torch.compiler.config.recompile_limit,
+        metavar="N",
+        help="maximum recompilations per compiled function before eager fallback",
+    )
+    group.add_argument(
+        "--compile-accumulated-recompile-limit",
+        type=int,
+        default=torch.compiler.config.accumulated_recompile_limit,
+        metavar="N",
+        help="maximum total recompilations across compiled functions",
+    )
+
+
+def add_checkpoint_args(parser: argparse.ArgumentParser, default_save_frequency: int = 1) -> None:
+    group = parser.add_argument_group("Checkpoint parameters")
+    group.add_argument(
+        "--save-frequency", type=int, default=default_save_frequency, metavar="N", help="frequency of model saving"
+    )
+    group.add_argument(
+        "--keep-last", type=int, metavar="N", help="number of recent checkpoints to keep (older ones are deleted)"
+    )
+    group.add_argument(
+        "--pretrained",
+        default=False,
+        action="store_true",
+        help="start with pretrained version of specified network (will download if not found locally)",
+    )
+    group.add_argument("--resume-epoch", type=int, metavar="N", help="epoch number to resume training from")
+    group.add_argument(
+        "--non-strict-weights",
+        default=False,
+        action="store_true",
+        help="allow non-strict loading of model weights (missing or unexpected keys in state_dict)",
+    )
+    group.add_argument(
+        "--load-states",
+        default=False,
+        action="store_true",
+        help="load optimizer, scheduler and scaler states when resuming",
+    )
+    group.add_argument("--load-scheduler", default=False, action="store_true", help="load only scheduler when resuming")
+
+
 def add_distributed_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Distributed training parameters")
     group.add_argument("--world-size", type=int, default=1, metavar="N", help="number of distributed processes")
@@ -558,5 +557,7 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--embed-dim must be positive")
     if args.context_length is not None and args.context_length <= 0:
         raise cli.ValidationError("--context-length must be positive")
+    if args.grad_accum_steps < 1:
+        raise cli.ValidationError("--grad-accum-steps must be >= 1")
     if args.model_ema_steps < 1:
         raise cli.ValidationError("--model-ema-steps must be >= 1")

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

@@ -0,0 +1,61 @@
+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 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],
+    *,
+    external_config: Optional[dict[str, Any]] = None,
+    **extra_states: Optional[dict[str, Any]],
+) -> None:
+    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,
+        )
+
+    if birder_training_utils.is_dist_available_and_initialized() is True:
+        dist.barrier()

birder_clip-0.0.2.dev4/birder_clip/data/datasets/webdataset.py

@@ -0,0 +1,106 @@
+import logging
+from collections.abc import Callable
+from functools import partial
+from typing import Any
+from typing import Optional
+
+import torch
+import webdataset as wds
+from birder.conf import settings
+from birder.data.datasets import webdataset as birder_wds
+
+from birder_clip.tokenizers import Tokenizer
+
+logger = logging.getLogger(__name__)
+
+
+def decode_caption(caption: Any, caption_json_key: str = "caption") -> str:
+    if isinstance(caption, dict):
+        if caption_json_key not in caption:
+            raise ValueError(f"WebDataset JSON sample missing '{caption_json_key}' key")
+
+        caption = caption[caption_json_key]
+
+    if isinstance(caption, bytes):
+        caption = caption.decode("utf-8")
+
+    if isinstance(caption, str) is False:
+        raise TypeError(f"WebDataset caption must be a string, got {type(caption).__name__}")
+
+    return caption  # type: ignore[no-any-return]
+
+
+def tokenize_caption(caption: str, tokenizer: Tokenizer) -> torch.Tensor:
+    return tokenizer([caption])[0]
+
+
+def make_wds_dataset(
+    wds_path: str | list[str],
+    dataset_size: int,
+    shuffle: bool,
+    samples_names: bool,
+    transform: Callable[..., torch.Tensor],
+    image_decoder: birder_wds.WDSImageDecoderSpec = "tv",
+    channels: int = settings.DEFAULT_NUM_CHANNELS,
+    tokenizer: Optional[Tokenizer] = None,
+    *,
+    caption_key: str = "txt;json",  # WebDataset picks the first present key, so txt takes precedence over json
+    caption_json_key: str = "caption",
+    cache_dir: Optional[str] = None,
+    shuffle_buffer_size: Optional[int] = None,
+    shuffle_initial_size: Optional[int] = None,
+) -> torch.utils.data.IterableDataset:
+    if shuffle is True:
+        shardshuffle = 500
+    else:
+        shardshuffle = False
+
+    dataset = wds.WebDataset(
+        wds_path, shardshuffle=shardshuffle, nodesplitter=wds.split_by_node, cache_dir=cache_dir, empty_check=False
+    )
+    if shuffle is True:
+        if shuffle_buffer_size is None:
+            shuffle_buffer_size = birder_wds.WDS_SHUFFLE_SIZE
+        if shuffle_initial_size is None:
+            shuffle_initial_size = birder_wds.WDS_INITIAL_SIZE
+
+        logger.debug(f"Using buffer size of {shuffle_buffer_size} for shuffle with {shuffle_initial_size} initial size")
+        dataset = dataset.shuffle(shuffle_buffer_size, initial=shuffle_initial_size)
+
+    return_keys = ["jpeg;jpg;png;webp"]
+    return_keys = return_keys + [caption_key]
+    if samples_names is True:
+        return_keys = ["__url__", "__key__"] + return_keys
+
+    if isinstance(image_decoder, str):
+        decoder = birder_wds.get_wds_image_decoder(image_decoder, channels)
+    else:
+        decoder = image_decoder
+
+    dataset = dataset.with_length(dataset_size, silent=True).decode(decoder).to_tuple(*return_keys)
+
+    caption_decoder = partial(decode_caption, caption_json_key=caption_json_key)
+    if samples_names is True:
+        dataset = dataset.map(birder_wds.decode_sample_name)
+        dataset = dataset.map_tuple(birder_wds.identity, transform, caption_decoder)
+    else:
+        dataset = dataset.map_tuple(transform, caption_decoder)
+
+    if tokenizer is not None:
+        text_transform = partial(tokenize_caption, tokenizer=tokenizer)
+        if samples_names is True:
+            dataset = dataset.map_tuple(birder_wds.identity, birder_wds.identity, text_transform)
+        else:
+            dataset = dataset.map_tuple(birder_wds.identity, text_transform)
+
+    return dataset
+
+
+def wds_size(wds_path: str, device: torch.device, select_suffix: str | tuple[str, ...] = ("txt", "json")) -> int:
+    return birder_wds.wds_size(wds_path, device, select_suffix=select_suffix)
+
+
+def prepare_wds_args(
+    data_path: str, size: Optional[int], device: torch.device, select_suffix: str | tuple[str, ...] = ("txt", "json")
+) -> tuple[str, int]:
+    return birder_wds.prepare_wds_args(data_path, size, device, select_suffix=select_suffix)

{birder_clip-0.0.2.dev3 → birder_clip-0.0.2.dev4}/birder_clip/loss/contrastive.py

@@ -1,6 +1,9 @@
 """
 CLIP loss, adapted from
 https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/loss.py
+
+Paper "Learning Transferable Visual Models From Natural Language Supervision",
+https://arxiv.org/abs/2103.00020
 """
 
 # Reference license: MIT
@@ -22,6 +25,14 @@ def gather_features(features: torch.Tensor) -> torch.Tensor:
 
 
 class CLIPLoss(torch.nn.Module):
+    """
+    CLIP symmetric contrastive loss
+
+    Implements the bidirectional InfoNCE objective from CLIP: image features
+    classify their matching text features, and text features classify their
+    matching image features, using the batch as negatives.
+    """
+
     def forward(
         self,
         image_features: torch.Tensor,

{birder_clip-0.0.2.dev3 → birder_clip-0.0.2.dev4}/birder_clip/net/clip.py

@@ -246,6 +246,10 @@ registry.register_model_config(
     },
 )
 
+
+# Weights
+####################
+
 registry.register_weights(
     "openai_clip_vit_l14",
     {
@@ -261,3 +265,18 @@ registry.register_weights(
         "net": {"network": "openai_clip_vit_l14"},
     },
 )
+registry.register_weights(
+    "pe_core_b16",
+    {
+        "description": "RoPEi ViT b16 image encoder pretrained by Meta FAIR using CLIP",
+        "resolution": (224, 224),
+        "context_length": 32,
+        "formats": {
+            "pt": {
+                "file_size": 1707.8,
+                "sha256": "11453d4a36fad6dbd802ec9fa35375ce0ae8b7b156a5ca45c0e87587df05290f",
+            }
+        },
+        "net": {"network": "pe_core_b16"},
+    },
+)

{birder_clip-0.0.2.dev3 → birder_clip-0.0.2.dev4}/birder_clip/net/text/base.py

@@ -1,4 +1,5 @@
 import copy
+import logging
 from typing import Any
 from typing import Optional
 
@@ -7,6 +8,8 @@ from torch import nn
 
 from birder_clip.model_registry import Task
 
+logger = logging.getLogger(__name__)
+
 
 class TextBaseNet(nn.Module):
     default_context_length = 77
@@ -38,4 +41,5 @@ class TextBaseNet(nn.Module):
         if new_context_length == self.context_length:
             return
 
+        logger.info(f"Adjusting model context length from {self.context_length} to {new_context_length}")
         self.context_length = new_context_length