birder-clip 0.0.2.dev6__tar.gz → 0.0.2.dev8__tar.gz

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: birder_clip
-Version: 0.0.2.dev6
+Version: 0.0.2.dev8
 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.6.0
+Requires-Dist: birder>=0.6.2
 Requires-Dist: ftfy>=6.3.1
 Requires-Dist: regex>=2025.7.29
 Requires-Dist: tqdm>=4.67.0
@@ -38,7 +38,7 @@ 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~=2026.1132; extra == "dev"
-Requires-Dist: coverage~=7.14.1; extra == "dev"
+Requires-Dist: coverage~=7.14.3; extra == "dev"
 Requires-Dist: debugpy; extra == "dev"
 Requires-Dist: flake8-pep585~=0.1.7; extra == "dev"
 Requires-Dist: flake8~=7.3.0; extra == "dev"
@@ -85,7 +85,7 @@ Full training is supported, but for large-scale CLIP pretraining you are probabl
 1. Ensure your environment meets the minimum requirements:
    - Python 3.11 or newer
    - PyTorch 2.10 or newer (installed for your hardware/driver stack)
-   - Birder 0.6.0 or newer
+   - Birder 0.6.2 or newer
 
 1. Install the latest Birder CLIP version:
 

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev8}/README.md

@@ -23,7 +23,7 @@ Full training is supported, but for large-scale CLIP pretraining you are probabl
 1. Ensure your environment meets the minimum requirements:
    - Python 3.11 or newer
    - PyTorch 2.10 or newer (installed for your hardware/driver stack)
-   - Birder 0.6.0 or newer
+   - Birder 0.6.2 or newer
 
 1. Install the latest Birder CLIP version:
 

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev8}/birder_clip/common/fs_ops.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Any
 from typing import NamedTuple
 from typing import Optional
+from typing import TypeAlias
 
 import torch
 from birder.common import cli
@@ -16,8 +17,10 @@ from birder.data.transforms.classification import inference_preset
 from birder_clip.common import lib
 from birder_clip.model_registry import Task
 from birder_clip.model_registry import registry
-from birder_clip.model_registry.manifest import EncoderMetadataType
 from birder_clip.model_registry.manifest import FileFormatType
+from birder_clip.model_registry.manifest import ImageEncoderMetadataType
+from birder_clip.model_registry.manifest import TextDecoderMetadataType
+from birder_clip.model_registry.manifest import TextEncoderMetadataType
 from birder_clip.net.base import BaseNet
 from birder_clip.net.base import SignatureType
 from birder_clip.tokenizers import Tokenizer
@@ -36,6 +39,8 @@ except ImportError:
 
 logger = logging.getLogger(__name__)
 
+ComponentMetadataType: TypeAlias = ImageEncoderMetadataType | TextEncoderMetadataType | TextDecoderMetadataType
+
 
 class ModelInfo(NamedTuple):
     signature: SignatureType
@@ -51,16 +56,18 @@ def write_config(network_name: str, net: BaseNet, signature: SignatureType, rgb_
         json.dump(model_config, handle, indent=2)
 
 
-def _split_encoder_metadata(encoder: Optional[EncoderMetadataType]) -> tuple[Optional[str], Optional[dict[str, Any]]]:
-    if encoder is None:
+def _split_component_metadata(
+    component: Optional[ComponentMetadataType],
+) -> tuple[Optional[str], Optional[dict[str, Any]]]:
+    if component is None:
         return (None, None)
-    if isinstance(encoder, str):
-        return (encoder, None)
+    if isinstance(component, str):
+        return (component, None)
 
-    if "network" not in encoder:
-        raise ValueError("Encoder metadata must include a 'network' field")
+    if "network" not in component:
+        raise ValueError("Component metadata must include a 'network' field")
 
-    return (None, encoder)  # type: ignore[return-value]
+    return (None, component)  # type: ignore[return-value]
 
 
 def model_path(
@@ -286,10 +293,12 @@ def load_checkpoint(
     tag: Optional[str] = None,
     image_encoder: Optional[str] = None,
     text_encoder: Optional[str] = None,
+    text_decoder: Optional[str] = None,
     embed_dim: Optional[int] = None,
     tokenizer: Optional[str] = None,
     image_encoder_config: Optional[dict[str, Any]] = None,
     text_encoder_config: Optional[dict[str, Any]] = None,
+    text_decoder_config: Optional[dict[str, Any]] = None,
     epoch: Optional[int] = None,
     new_size: Optional[tuple[int, int]] = None,
     new_context_length: Optional[int] = None,
@@ -300,6 +309,7 @@ def load_checkpoint(
         tag=tag,
         image_encoder=image_encoder,
         text_encoder=text_encoder,
+        text_decoder=text_decoder,
         embed_dim=embed_dim,
         tokenizer=tokenizer,
     )
@@ -329,10 +339,12 @@ def load_checkpoint(
         checkpoint_config,
         image_encoder=image_encoder,
         text_encoder=text_encoder,
+        text_decoder=text_decoder,
         embed_dim=embed_dim,
         tokenizer=tokenizer,
         image_encoder_config=image_encoder_config,
         text_encoder_config=text_encoder_config,
+        text_decoder_config=text_decoder_config,
         input_channels=input_channels,
         image_size=size,
         context_length=context_length,
@@ -364,10 +376,12 @@ def load_model(
     tag: Optional[str] = None,
     image_encoder: Optional[str] = None,
     text_encoder: Optional[str] = None,
+    text_decoder: Optional[str] = None,
     embed_dim: Optional[int] = None,
     tokenizer: Optional[str] = None,
     image_encoder_config: Optional[dict[str, Any]] = None,
     text_encoder_config: Optional[dict[str, Any]] = None,
+    text_decoder_config: Optional[dict[str, Any]] = None,
     epoch: Optional[int] = None,
     new_size: Optional[tuple[int, int]] = None,
     new_context_length: Optional[int] = None,
@@ -381,6 +395,7 @@ def load_model(
             tag=tag,
             image_encoder=image_encoder,
             text_encoder=text_encoder,
+            text_decoder=text_decoder,
             embed_dim=embed_dim,
             tokenizer=tokenizer,
         )
@@ -420,10 +435,12 @@ def load_model(
         checkpoint_config,
         image_encoder=image_encoder,
         text_encoder=text_encoder,
+        text_decoder=text_decoder,
         embed_dim=embed_dim,
         tokenizer=tokenizer,
         image_encoder_config=image_encoder_config,
         text_encoder_config=text_encoder_config,
+        text_decoder_config=text_decoder_config,
         input_channels=input_channels,
         image_size=size,
         context_length=context_length,
@@ -525,15 +542,17 @@ def load_pretrained_model(
     if model_metadata["task"] != Task.IMAGE_TEXT:
         raise ValueError(f"Unknown model type: {model_metadata['task']}")
 
-    image_encoder, image_config = _split_encoder_metadata(model_metadata["net"].get("image_encoder", None))
-    text_encoder, text_config = _split_encoder_metadata(model_metadata["net"].get("text_encoder", None))
+    image_encoder, image_config = _split_component_metadata(model_metadata["net"].get("image_encoder", None))
+    text_encoder, text_config = _split_component_metadata(model_metadata["net"].get("text_encoder", None))
+    text_decoder, decoder_config = _split_component_metadata(model_metadata["net"].get("text_decoder", None))
 
     pretrained_config: dict[str, Any] = {}
     if image_config is not None:
         pretrained_config["image"] = image_config
-
     if text_config is not None:
         pretrained_config["text"] = text_config
+    if decoder_config is not None:
+        pretrained_config["decoder"] = decoder_config
 
     if custom_config is not None:
         pretrained_config.update(custom_config)
@@ -550,6 +569,7 @@ def load_pretrained_model(
         tag=model_metadata["net"].get("tag", None),
         image_encoder=image_encoder,
         text_encoder=text_encoder,
+        text_decoder=text_decoder,
         embed_dim=model_metadata["net"].get("embed_dim", None),
         tokenizer=model_metadata["net"].get("tokenizer", None),
         inference=inference,

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev8}/birder_clip/common/lib.py

@@ -3,14 +3,17 @@ from typing import Any
 from typing import Optional
 
 from birder.data.transforms.classification import RGBType
+from birder.version import __version__ as birder_version
 
 from birder_clip.conf import settings
 from birder_clip.model_registry import registry
 from birder_clip.net.base import BaseNet
 from birder_clip.net.base import SignatureType
-from birder_clip.version import __version__
+from birder_clip.version import __version__ as birder_clip_version
 
-MODEL_CONFIG_RESERVED_KEYS = frozenset({"image", "text", "tokenizer", "embed_dim", "embed-dim"})
+MODEL_CONFIG_RESERVED_KEYS = frozenset(
+    {"image", "text", "decoder", "tokenizer", "embed_dim", "embed-dim", "keep_ratio"}
+)
 
 
 def get_size_from_signature(signature: SignatureType) -> tuple[int, int]:
@@ -37,6 +40,7 @@ def get_image_text_network_name(
     tag: Optional[str] = None,
     image_encoder: Optional[str] = None,
     text_encoder: Optional[str] = None,
+    text_decoder: Optional[str] = None,
     embed_dim: Optional[int] = None,
     tokenizer: Optional[str] = None,
 ) -> str:
@@ -45,6 +49,8 @@ def get_image_text_network_name(
         parts.append(image_encoder)
     if text_encoder is not None and text_encoder != "transformer_encoder":
         parts.append(text_encoder)
+    if text_decoder is not None and text_decoder != "conditioned_decoder":
+        parts.append(text_decoder)
 
     if registry.exists(network) is True:
         default_tokenizer = registry.get_default_tokenizer(network)
@@ -71,10 +77,12 @@ def get_image_text_model_config(
     *,
     image_encoder: Optional[str] = None,
     text_encoder: Optional[str] = None,
+    text_decoder: Optional[str] = None,
     embed_dim: Optional[int] = None,
     tokenizer: Optional[str] = None,
     image_encoder_config: Optional[dict[str, Any]] = None,
     text_encoder_config: Optional[dict[str, Any]] = None,
+    text_decoder_config: Optional[dict[str, Any]] = None,
     input_channels: Optional[int] = None,
     image_size: Optional[tuple[int, int]] = None,
     context_length: Optional[int] = None,
@@ -86,7 +94,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):
+            if key in {"image", "text", "decoder"} and isinstance(value, dict):
                 model_config[key] = {**model_config.get(key, {}), **value}
             else:
                 model_config[key] = value
@@ -111,7 +119,7 @@ def get_image_text_model_config(
 
         model_config["image"] = image_config
 
-    if text_encoder is not None or text_encoder_config is not None or context_length is not None:
+    if text_encoder is not None or text_encoder_config is not None or "text" in model_config:
         text_config = model_config.get("text", {}).copy()
         if text_encoder is not None:
             # String encoder metadata replaces only the encoder name.
@@ -124,6 +132,17 @@ def get_image_text_model_config(
 
         model_config["text"] = text_config
 
+    if text_decoder is not None or text_decoder_config is not None or "decoder" in model_config:
+        decoder_config = model_config.get("decoder", {}).copy()
+        if text_decoder is not None:
+            decoder_config["network"] = text_decoder
+        if text_decoder_config is not None:
+            decoder_config["config"] = {**decoder_config.get("config", {}), **text_decoder_config}
+        if context_length is not None:
+            decoder_config["context_length"] = context_length
+
+        model_config["decoder"] = decoder_config
+
     if embed_dim is not None:
         model_config["embed_dim"] = embed_dim
     if tokenizer is not None:
@@ -143,7 +162,8 @@ def get_image_text_network_config(net: BaseNet, signature: SignatureType, rgb_st
         model_config = net.config
 
     return {
-        "birder_clip_version": __version__,
+        "birder_clip_version": birder_clip_version,
+        "birder_version": birder_version,
         "name": model_name,
         "registered_name": registered_name,
         "task": net.task,

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev8}/birder_clip/common/training_cli.py

@@ -27,6 +27,7 @@ def add_model_args(parser: argparse.ArgumentParser) -> None:
         help="pretrained Birder image model weights path to load into the image encoder",
     )
     parser.add_argument("--text-encoder", type=str, help="the text encoder to use")
+    parser.add_argument("--text-decoder", type=str, help="the text decoder to use")
     parser.add_argument("--embed-dim", type=int, metavar="N", help="shared image-text embedding dimension")
     parser.add_argument("--tokenizer", type=str, help="the tokenizer to use")
     parser.add_argument(
@@ -44,11 +45,21 @@ def add_model_args(parser: argparse.ArgumentParser) -> None:
         action=cli.FlexibleDictAction,
         help="override the text encoder configuration, accepts key-value pairs or JSON",
     )
+    parser.add_argument(
+        "--text-decoder-config",
+        action=cli.FlexibleDictAction,
+        help="override the text decoder configuration, accepts key-value pairs or JSON",
+    )
+    parser.add_argument(
+        "--openvision-v2-keep-ratio", type=float, help="OpenVision v2 image token keep ratio for caption decoding"
+    )
 
 
 def add_loss_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Loss parameters")
-    group.add_argument("--loss", type=str, choices=["clip", "coca"], default="clip", help="loss function to use")
+    group.add_argument(
+        "--loss", type=str, choices=["clip", "coca", "caption"], default="clip", help="loss function to use"
+    )
     group.add_argument(
         "--coca-caption-loss-weight", type=float, default=1.0, help="weight assigned to CoCa caption loss"
     )
@@ -662,13 +673,16 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--grad-accum-steps must be >= 1")
     if args.grad_accum_cache_negatives is True and args.grad_accum_steps == 1:
         raise cli.ValidationError("--grad-accum-cache-negatives requires --grad-accum-steps greater than 1")
-    if args.grad_accum_cache_negatives is True and args.loss == "coca":
+    if args.grad_accum_cache_negatives is True and args.loss != "clip":
         raise cli.ValidationError("--grad-accum-cache-negatives is only supported with --loss clip")
 
     if args.coca_caption_loss_weight < 0.0:
         raise cli.ValidationError("--coca-caption-loss-weight must be non-negative")
     if args.coca_contrastive_loss_weight < 0.0:
         raise cli.ValidationError("--coca-contrastive-loss-weight must be non-negative")
+    if args.openvision_v2_keep_ratio is not None:
+        if args.openvision_v2_keep_ratio <= 0.0 or args.openvision_v2_keep_ratio > 1.0:
+            raise cli.ValidationError("--openvision-v2-keep-ratio must be in range of (0, 1]")
 
     # EMA
     if args.model_ema_steps < 1:

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev8}/birder_clip/common/training_utils.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import logging
 from pathlib import Path
 from typing import Any
@@ -11,6 +12,7 @@ from birder.common import training_utils as birder_training_utils
 
 from birder_clip.common import fs_ops
 from birder_clip.conf import settings
+from birder_clip.version import __version__ as birder_clip_version
 
 
 def setup_file_logging(log_file_path: str | Path) -> logging.Handler:
@@ -28,6 +30,18 @@ def setup_file_logging(log_file_path: str | Path) -> logging.Handler:
     return file_handler
 
 
+def make_training_args_payload(args: argparse.Namespace) -> dict[str, Any]:
+    return {
+        "birder_clip_version": birder_clip_version,
+        **birder_training_utils.make_training_args_payload(args),
+    }
+
+
+def write_training_args_json(path: Path, args: argparse.Namespace) -> None:
+    with open(path.joinpath("training_args.json"), "w", encoding="utf-8") as handle:
+        json.dump(make_training_args_payload(args), handle, indent=2)
+
+
 def save_training_checkpoint(
     args: argparse.Namespace,
     network_name: str,

birder_clip-0.0.2.dev8/birder_clip/inference/data_parallel.py

@@ -0,0 +1,118 @@
+"""
+Inference-optimized multi-GPU parallelization for image-text models
+
+This module provides ZeroShotInferenceDataParallel, a CLIP-style zero-shot
+specialization of Birder's InferenceDataParallel.
+"""
+
+from typing import Optional
+
+import torch
+from birder.inference.data_parallel import InferenceDataParallel
+
+from birder_clip.inference.zero_shot import ZeroShotInference
+from birder_clip.net.base import BaseNet
+
+
+class ZeroShotInferenceDataParallel(InferenceDataParallel):
+    """
+    Distributes zero-shot image inference batches across multiple GPUs
+
+    This wrapper scatters the image batch across devices and keeps a replicated
+    copy of the zero-shot text embeddings on each device. Each replica computes
+    image embeddings and zero-shot logits locally before outputs are gathered.
+
+    Important
+    ---------
+    This class assumes the model is already configured for inference mode
+    (i.e., loaded with inference=True in load_model or manually set to eval mode
+    with requires_grad=False on all parameters).
+    """
+
+    def __init__(
+        self,
+        module: BaseNet,
+        text_embeddings: torch.Tensor,
+        device_ids: Optional[list[int]] = None,
+        output_device: Optional[int | str | torch.device] = None,
+        compile_replicas: bool = False,
+        compile_methods: Optional[list[str]] = None,
+        compile_mode: Optional[str] = None,
+    ) -> None:
+        if compile_methods is None:
+            compile_methods = ["encode_image", "forward_logits"]
+
+        super().__init__(
+            module,
+            device_ids=device_ids,
+            output_device=output_device,
+            compile_replicas=compile_replicas,
+            compile_methods=compile_methods,
+            compile_mode=compile_mode,
+        )
+        self.set_text_embeddings(text_embeddings)
+
+    def set_text_embeddings(self, text_embeddings: torch.Tensor) -> None:
+        if text_embeddings.ndim != 2:
+            raise ValueError(
+                f"text_embeddings must be a 2D tensor of shape (num_classes, embedding_size), "
+                f"got shape {text_embeddings.size()}"
+            )
+
+        self.text_embeddings = [
+            text_embeddings.to(f"cuda:{device_id}", non_blocking=True) for device_id in self.device_ids
+        ]
+        self.inference_modules = [
+            ZeroShotInference(replica, embeddings) for replica, embeddings in zip(self.replicas, self.text_embeddings)
+        ]
+
+    def forward(  # type: ignore[override] # pylint: disable=arguments-differ
+        self, inputs: torch.Tensor, *, tta: bool = False, return_logits: bool = False
+    ) -> torch.Tensor:
+        """
+        Run zero-shot inference distributed across GPUs
+
+        Parameters
+        ----------
+        inputs
+            Input image batch to process.
+        tta
+            Run inference with oversampling.
+        return_logits
+            If True, return raw logits instead of probabilities after softmax.
+        """
+
+        if len(self.device_ids) == 1:
+            output = self.inference_modules[0](
+                inputs,
+                tta=tta,
+                return_logits=return_logits,
+            )
+            return self._gather([output])
+
+        scattered = self._scatter(inputs, {})
+
+        outputs = []
+        for inference, (input_chunk, _), device_id in zip(self.inference_modules, scattered, self.device_ids):
+            if input_chunk is not None and input_chunk.size(0) > 0:
+                with torch.cuda.device(device_id):
+                    output = inference(
+                        input_chunk,
+                        tta=tta,
+                        return_logits=return_logits,
+                    )
+                    outputs.append(output)
+            else:
+                outputs.append(None)
+
+        return self._gather(outputs)
+
+    def __repr__(self) -> str:
+        return (
+            f"ZeroShotInferenceDataParallel(\n"
+            f"  devices={self.device_ids},\n"
+            f"  output_device={self.output_device},\n"
+            f"  src_device={self.src_device},\n"
+            f"  text_embeddings_shape={tuple(self.text_embeddings[0].shape)}\n"
+            f")"
+        )

birder_clip-0.0.2.dev8/birder_clip/inference/image_embeddings.py

@@ -0,0 +1,63 @@
+import sys
+from collections.abc import Iterator
+from typing import Optional
+
+import numpy as np
+import numpy.typing as npt
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from birder_clip.net.base import BaseNet
+
+DataloaderInferenceResult = tuple[list[str], npt.NDArray[np.float32]]
+
+
+def infer_dataloader_iter(
+    device: torch.device,
+    net: BaseNet,
+    dataloader: DataLoader,
+    model_dtype: torch.dtype = torch.float32,
+    amp: bool = False,
+    amp_dtype: Optional[torch.dtype] = None,
+    num_samples: Optional[int] = None,
+    chunk_size: Optional[float] = None,
+) -> Iterator[DataloaderInferenceResult]:
+    if chunk_size is None:
+        chunk_size = float("inf")
+
+    net.to(device, dtype=model_dtype)
+    embeddings_list: list[npt.NDArray[np.float32]] = []
+    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):
+                embeddings = net.encode_image(inputs, normalize=True)
+                embeddings = embeddings.cpu().float().numpy()
+
+            embeddings_list.append(embeddings)
+
+            # Set sample list
+            sample_paths.extend(file_paths)
+
+            # 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(embeddings_list, axis=0))
+
+                # Reset for next chunk
+                embeddings_list = []
+                sample_paths = []
+                sample_count = 0
+
+    if len(embeddings_list) > 0:
+        yield (sample_paths, np.concatenate(embeddings_list, axis=0))

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev8}/birder_clip/inference/zero_shot.py

@@ -13,6 +13,7 @@ from collections.abc import Callable
 from collections.abc import Iterator
 from collections.abc import Sequence
 from typing import Optional
+from typing import Protocol
 
 import numpy as np
 import numpy.typing as npt
@@ -27,6 +28,42 @@ from birder_clip.net.base import BaseNet
 from birder_clip.tokenizers.base import Tokenizer
 
 
+class ZeroShotInferenceModule(Protocol):
+    def __call__(self, inputs: torch.Tensor, *, tta: bool = False, return_logits: bool = False) -> torch.Tensor: ...
+
+
+class ZeroShotInference:
+    def __init__(self, net: BaseNet, text_embeddings: torch.Tensor) -> None:
+        self.net = net
+        self.text_embeddings = text_embeddings
+
+    def __call__(self, inputs: torch.Tensor, *, tta: bool = False, return_logits: bool = False) -> torch.Tensor:
+        inputs = inputs.to(self.text_embeddings.device, non_blocking=True)
+        if tta is True:
+            _, _, H, W = inputs.size()
+            crop_h = int(H * 0.8)
+            crop_w = int(W * 0.8)
+            tta_inputs = five_crop(inputs, size=[crop_h, crop_w])
+            t = v2.Resize((H, W), interpolation=v2.InterpolationMode.BICUBIC, antialias=True)
+            outs = []
+            for tta_input in tta_inputs:
+                image_embeddings = self.net.encode_image(t(tta_input), normalize=True)
+                logits = self.net.forward_logits(image_embeddings, self.text_embeddings)
+                if return_logits is True:
+                    outs.append(logits)
+                else:
+                    outs.append(F.softmax(logits, dim=-1))
+
+            return torch.stack(outs).mean(dim=0)
+
+        image_embeddings = self.net.encode_image(inputs, normalize=True)
+        logits = self.net.forward_logits(image_embeddings, self.text_embeddings)
+        if return_logits is True:
+            return logits
+
+        return F.softmax(logits, dim=-1)
+
+
 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]
 
@@ -66,39 +103,10 @@ def build_class_text_embeddings(
 DataloaderInferenceResult = tuple[list[str], npt.NDArray[np.float32], npt.NDArray[np.int64]]
 
 
-def infer_batch(
-    net: BaseNet, inputs: torch.Tensor, text_embeddings: torch.Tensor, tta: bool = False, return_logits: bool = False
-) -> torch.Tensor:
-    if tta is True:
-        _, _, H, W = inputs.size()
-        crop_h = int(H * 0.8)
-        crop_w = int(W * 0.8)
-        tta_inputs = five_crop(inputs, size=[crop_h, crop_w])
-        t = v2.Resize((H, W), interpolation=v2.InterpolationMode.BICUBIC, antialias=True)
-        outs = []
-        for tta_input in tta_inputs:
-            image_embeddings = net.encode_image(t(tta_input), normalize=True)
-            logits = net.forward_logits(image_embeddings, text_embeddings)
-            if return_logits is True:
-                outs.append(logits)
-            else:
-                outs.append(F.softmax(logits, dim=-1))
-
-        return torch.stack(outs).mean(dim=0)
-
-    image_embeddings = net.encode_image(inputs, normalize=True)
-    logits = net.forward_logits(image_embeddings, text_embeddings)
-    if return_logits is True:
-        return logits
-
-    return F.softmax(logits, dim=-1)
-
-
 def infer_dataloader_iter(
     device: torch.device,
-    net: BaseNet,
+    net: ZeroShotInferenceModule,
     dataloader: DataLoader,
-    text_embeddings: torch.Tensor,
     tta: bool = False,
     return_logits: bool = False,
     model_dtype: torch.dtype = torch.float32,
@@ -111,7 +119,6 @@ def infer_dataloader_iter(
     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] = []
@@ -121,9 +128,9 @@ def infer_dataloader_iter(
             batch_size = inputs.size(0)
 
             # Inference
-            inputs = inputs.to(device, dtype=model_dtype)
+            inputs = inputs.to(dtype=model_dtype)
             with torch.amp.autocast(device.type, enabled=amp, dtype=amp_dtype):
-                out = infer_batch(net, inputs, text_embeddings, return_logits=return_logits, tta=tta)
+                out = net(inputs, return_logits=return_logits, tta=tta)
                 out = out.cpu().float().numpy()
 
             out_list.append(out)

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev8}/birder_clip/loss/__init__.py

@@ -1,7 +1,9 @@
+from birder_clip.loss.caption import CaptionLoss
 from birder_clip.loss.coca import CoCaLoss
 from birder_clip.loss.contrastive import CLIPLoss
 
 __all__ = [
+    "CaptionLoss",
     "CoCaLoss",
     "CLIPLoss",
 ]