visbench 0.2.0__py3-none-any.whl

visbench/backbones/clip.py

@@ -0,0 +1,247 @@
+"""OpenCLIP ViT-B image-tower backbone — v0.1.
+
+Only the visual tower is exposed; the text tower is out of scope for probing.
+Note that CLIP uses its own normalisation constants, not ImageNet's.
+"""
+
+import warnings
+
+import torch
+from PIL import Image
+from torchvision import transforms
+
+from visbench.backbones.base import BaseBackbone
+from visbench.registry import register_backbone
+from visbench.types import LayerOutput
+from visbench.utils.image import CLIP_MEAN, CLIP_STD
+
+__all__ = ["CLIP"]
+
+#: Registered name -> (open_clip model name, pretrained tag, embed dim, patch size).
+#:
+#: The ``-quickgelu`` suffix is not cosmetic. OpenAI's original CLIP weights
+#: were trained with QuickGELU, and open_clip pairs them with a plain-GELU
+#: architecture unless asked otherwise — it warns and continues, producing a
+#: model that loads cleanly and computes subtly wrong activations. The pairing
+#: is validated in ``__init__`` so that combination raises instead.
+_VARIANTS = {
+    "clip_vitb16": ("ViT-B-16-quickgelu", "openai", 768, 16),
+    "clip_vitb32": ("ViT-B-32-quickgelu", "openai", 768, 32),
+}
+
+#: Dim of the shared image-text embedding for ViT-B, after ``visual.proj``.
+_PROJECTED_DIM = 512
+
+#: Substring identifying open_clip's QuickGELU warnings, matched case-insensitively.
+#:
+#: Deliberately one token rather than a sentence. This guard was originally a
+#: ``warnings.filterwarnings("error", message=".*QuickGELU mismatch.*")``, and
+#: open_clip has never emitted the phrase "QuickGELU mismatch" — so the filter
+#: never matched, the warning was never promoted, and the guard was dead code
+#: from the day it was written. Its own test only caught it under ``-m slow``,
+#: which CI does not run.
+#:
+#: open_clip warns in *both* directions — weights trained with QuickGELU loaded
+#: under a plain-GELU config, and the reverse — with different wording each way
+#: (see ``open_clip/factory.py``). Both are genuine mismatches. This is the one
+#: token common to the two, and the only part of the wording worth depending on.
+_QUICKGELU_MARKER = "quickgelu"
+
+
+def _promote_quickgelu_warning(
+    caught: "list[warnings.WarningMessage]", model_name: str, pretrained: str
+) -> None:
+    """Turn open_clip's QuickGELU warning into an error; re-emit anything else.
+
+    open_clip only warns on an activation mismatch and hands back a model that
+    loads cleanly and computes subtly wrong features. A warning in the middle of
+    a benchmark run is a warning nobody reads, and this is the one failure mode
+    here that silently changes a number rather than raising.
+
+    Every other warning is re-emitted rather than swallowed: recording warnings
+    suppresses them, and a guard against one specific problem has no business
+    hiding unrelated deprecation notices from the caller.
+    """
+    mismatch: str | None = None
+    for entry in caught:
+        text = str(entry.message)
+        if mismatch is None and _QUICKGELU_MARKER in text.lower():
+            mismatch = text
+            continue
+        # stacklevel=2 points at CLIP.__init__ rather than this helper. The true
+        # origin inside open_clip is not recoverable once a warning is recorded.
+        warnings.warn(entry.message, stacklevel=2)
+
+    if mismatch is not None:
+        raise RuntimeError(
+            f"{model_name} + {pretrained!r} is a QuickGELU mismatch: {mismatch} "
+            "That combination loads cleanly and computes wrong activations."
+        )
+
+
+@register_backbone("clip_vitb16", variant="clip_vitb16")
+@register_backbone("clip_vitb32", variant="clip_vitb32")
+class CLIP(BaseBackbone):
+    """CLIP visual tower.
+
+    Has a CLS token, so default pooling is CLS — but note this is the
+    pre-projection CLS, i.e. the representation used for probing, not the
+    projected image embedding used for text alignment. Which of the two a task
+    wants is a real decision; document it wherever it is made.
+
+    **Pre-projection is the default here.** CLIP's visual tower ends with a
+    linear projection into the space shared with the text encoder, and
+    ``encode_image`` returns that 512-d vector. VisBench returns the 768-d CLS
+    token from *before* it, because the projection is trained to discard
+    whatever does not help match a caption — which is exactly the visual detail
+    a mid-level probe exists to measure. It is also the wrong comparison to
+    draw against DINOv2, which has no such head. Pass ``use_projection=True``
+    for the 512-d embedding, which is what a published zero-shot number used.
+
+    Patch size is 16 or 32 against DINOv2's 14, so the two produce different
+    grids at the same input resolution. That is why correspondence measures
+    error in patch widths rather than pixels.
+    """
+
+    has_cls_token = True
+
+    def __init__(
+        self,
+        variant: str = "clip_vitb16",
+        device: str | None = None,
+        image_size: int = 224,
+        use_projection: bool = False,
+    ) -> None:
+        """Load the OpenCLIP visual tower for ``variant``, freeze it, set eval mode."""
+        super().__init__(device)
+
+        if variant not in _VARIANTS:
+            raise ValueError(
+                f"Unknown CLIP variant {variant!r}; expected one of {sorted(_VARIANTS)}"
+            )
+        model_name, pretrained, width, patch_size = _VARIANTS[variant]
+
+        if image_size % patch_size != 0:
+            raise ValueError(
+                f"image_size={image_size} is not a multiple of the patch size {patch_size}. "
+                "A ragged final patch would silently change the grid shape."
+            )
+
+        self.name = variant
+        self.variant = variant
+        self.model_name = model_name
+        self.pretrained = pretrained
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.use_projection = use_projection
+        self.embed_dim = _PROJECTED_DIM if use_projection else width
+
+        try:
+            import open_clip
+        except ImportError as exc:  # pragma: no cover - needs the extra uninstalled
+            raise ImportError(
+                "The CLIP backbone needs open_clip_torch. Install it with "
+                "`pip install visbench[clip]`."
+            ) from exc
+
+        # Record rather than filter-to-error: open_clip's wording is the only
+        # signal available, and matching it as a regex is what broke this guard
+        # once. See _promote_quickgelu_warning.
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            model = open_clip.create_model(model_name, pretrained=pretrained)
+        _promote_quickgelu_warning(caught, model_name, pretrained)
+
+        self.model = model.visual
+        self._transform = transforms.Compose(
+            [
+                transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
+                transforms.CenterCrop(image_size),
+                transforms.ToTensor(),
+                # CLIP's own constants, not ImageNet's — see visbench.utils.image.
+                transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
+            ]
+        )
+        self._finalize()
+
+    @property
+    def num_layers(self) -> int:
+        """Residual attention blocks in the visual transformer."""
+        return len(self.model.transformer.resblocks)
+
+    def _forward_features(
+        self,
+        image: torch.Tensor,
+        layers: list[int],
+    ) -> list[LayerOutput]:
+        """Run the visual transformer once, one output per requested layer.
+
+        Uses ``forward_intermediates``, open_clip's counterpart to DINOv2's
+        ``get_intermediate_layers``, which takes the whole index list.
+        """
+        _, _, height, width = image.shape
+        patch = self.patch_size
+        assert patch is not None  # set in __init__; Optional only on the base class
+        if height % patch or width % patch:
+            raise ValueError(f"Input {height}x{width} is not a multiple of patch size {patch}")
+        grid_hw = (height // patch, width // patch)
+
+        outputs = self.model.forward_intermediates(
+            image,
+            indices=list(layers),
+            output_fmt="NLC",
+            output_extra_tokens=True,
+            intermediates_only=True,
+        )
+
+        expected = grid_hw[0] * grid_hw[1]
+        result: list[LayerOutput] = []
+        for position, index in enumerate(layers):
+            patch_tokens = outputs["image_intermediates"][position]
+            cls_token = outputs["image_intermediates_prefix"][position][:, 0]
+
+            # forward_intermediates returns raw block outputs, so the final
+            # LayerNorm has not been applied. DINOv2's equivalent normalises by
+            # default; applying it here keeps "the features" meaning the same
+            # thing across backbones, and it is what `proj` expects downstream —
+            # ln_post(cls) @ proj reproduces encode_image exactly.
+            #
+            # Applied to every layer, not just the last. ln_post is trained for
+            # the final block's scale, so on an intermediate layer it is a
+            # convention rather than a reconstruction — but an unnormalised
+            # layer sitting next to normalised ones in the same pyramid would
+            # hand a multiscale head stages whose magnitudes differ by a factor
+            # the head would have to unlearn.
+            patch_tokens = self.model.ln_post(patch_tokens)
+            cls_token = self.model.ln_post(cls_token)
+
+            if self.use_projection:
+                if self.model.proj is None:
+                    raise RuntimeError(f"{self.model_name} has no visual projection")
+                patch_tokens = patch_tokens @ self.model.proj
+                cls_token = cls_token @ self.model.proj
+
+            if patch_tokens.shape[1] != expected:
+                raise RuntimeError(
+                    f"CLIP layer {index} returned {patch_tokens.shape[1]} patch tokens, "
+                    f"expected {expected} for a {grid_hw[0]}x{grid_hw[1]} grid"
+                )
+            result.append((patch_tokens, cls_token, grid_hw))
+        return result
+
+    def preprocess(self, images: Image.Image | list) -> torch.Tensor:
+        """Resize to 224 and apply CLIP normalisation constants."""
+        if isinstance(images, Image.Image):
+            images = [images]
+        return torch.stack([self._transform(img.convert("RGB")) for img in images])
+
+    def cache_key(self) -> str:
+        """``"clip/<model>/<pretrained-tag>/<resolution>/<head>"``.
+
+        The pretrained tag is in the key because ``openai`` and ``laion2b``
+        weights are different models behind one name, and ``head`` because the
+        projected and pre-projection vectors are different representations of
+        the same forward pass.
+        """
+        head = "proj" if self.use_projection else "preproj"
+        return f"clip/{self.model_name}/{self.pretrained}/{self.image_size}/{head}"

visbench/backbones/custom.py

@@ -0,0 +1,311 @@
+"""Wrap an arbitrary ``nn.Module`` as a VisBench backbone. v0.2.
+
+The escape hatch: any model a user already has — a fine-tuned checkpoint, an
+architecture VisBench has never heard of, something from a paper's repo — probed
+by the same tasks as DINOv2 and CLIP, without adding a module to this package.
+
+Constructed directly rather than looked up by name, since a registry name
+cannot carry an ``nn.Module``::
+
+    backbone = CustomBackbone(my_model, preprocess=my_transform, name="mine")
+    visbench.run(backbone, "retrieval", dataset)
+
+To give a custom backbone a registry name of its own, subclass
+:class:`~visbench.backbones.base.BaseBackbone` and apply
+:func:`visbench.register_backbone` — that path is unchanged and is how the
+built-in backbones work.
+"""
+
+import hashlib
+from collections.abc import Callable
+from typing import Any
+
+import torch
+import torch.nn as nn
+from PIL import Image
+
+from visbench.backbones.base import BaseBackbone
+from visbench.types import LayerOutput
+
+__all__ = ["CustomBackbone"]
+
+#: Return signature of a user-supplied feature function.
+FeatureFn = Callable[[nn.Module, torch.Tensor], tuple]
+
+#: Multi-layer counterpart: ``(module, image_batch, layers) -> list`` of
+#: ``(patch_tokens, cls_token, grid_hw)``, one per requested index. Kept
+#: separate from :data:`FeatureFn` rather than overloading its signature, so a
+#: function written for one convention can never be called under the other.
+LayerFeatureFn = Callable[[nn.Module, torch.Tensor, list], list]
+
+
+def hash_weights(module: nn.Module, sample_bytes: int = 1 << 20) -> str:
+    """Short hash of a module's parameters, for the cache key.
+
+    A custom backbone has no upstream commit or pretrained tag to point at, so
+    the weights themselves are the only honest identifier. Hashing them means a
+    fine-tuned checkpoint automatically gets a different cache key from the one
+    it was fine-tuned from — the alternative is a user-supplied string that is
+    correct only if they remember to change it.
+
+    Large tensors are sampled rather than read whole: a full pass over a
+    multi-GB state dict on every construction would cost more than the forward
+    passes it protects. Shapes and dtypes are always folded in, so a change in
+    architecture is caught even when the sampled bytes collide.
+    """
+    digest = hashlib.sha256()
+    for name, tensor in sorted(module.state_dict().items()):
+        flat = tensor.detach().cpu().flatten()
+        digest.update(f"{name}|{tuple(tensor.shape)}|{tensor.dtype}".encode())
+        if flat.numel() == 0:
+            continue
+        values = flat.to(torch.float64)
+        # A cheap moment-based summary plus a byte sample: catches retraining
+        # and fine-tuning, which is what this needs to detect.
+        digest.update(f"{values.sum().item():.6e}|{values.abs().max().item():.6e}".encode())
+        digest.update(flat[: sample_bytes // flat.element_size()].numpy().tobytes())
+    return digest.hexdigest()[:16]
+
+
+class CustomBackbone(BaseBackbone):
+    """Adapt a user-supplied ``nn.Module`` to the VisBench backbone contract.
+
+    Parameters
+    ----------
+    module:
+        Any ``nn.Module``. It is frozen and set to eval mode — probing measures
+        fixed representations, and unfreezing is v0.3 scope.
+    preprocess:
+        Callable turning one PIL image into a ``(3, H, W)`` tensor, typically a
+        torchvision ``Compose``. Required, and deliberately so: normalisation
+        constants cannot be guessed, and getting them wrong is a silent
+        accuracy loss rather than an error.
+    feature_fn:
+        Optional ``(module, image_batch) -> (patch_tokens, cls_token, grid_hw)``
+        for full control. Without it, the module's own ``forward`` is called and
+        its output interpreted — see :meth:`_forward_features`.
+    has_cls_token:
+        Whether the token sequence starts with a CLS token. Only consulted when
+        the module returns tokens; ignored for a conv map.
+    layer_feature_fn:
+        Optional ``(module, image_batch, layers) -> [(tokens, cls, grid_hw), ...]``
+        enabling multi-layer extraction. Requires ``num_layers``. VisBench
+        cannot tap the intermediate activations of an arbitrary module — there
+        is no equivalent of ``get_intermediate_layers`` to call — so this is
+        the seam where a user says how their own model exposes depth.
+    num_layers:
+        How many depths ``layer_feature_fn`` can serve. Left at 1, the module
+        exposes only its final output and any multi-layer request is rejected,
+        which is the right default: silently returning the same map several
+        times would let a multiscale head report a single-layer result.
+    weights_id:
+        Identifier for these weights in the cache key. Defaults to a hash of
+        the module's parameters, which is usually what you want.
+    """
+
+    def __init__(
+        self,
+        module: nn.Module,
+        preprocess: Callable[[Image.Image], torch.Tensor],
+        name: str = "custom",
+        embed_dim: int | None = None,
+        has_cls_token: bool = False,
+        patch_size: int | None = None,
+        feature_fn: FeatureFn | None = None,
+        weights_id: str | None = None,
+        image_size: int = 224,
+        device: str | None = None,
+        layer_feature_fn: LayerFeatureFn | None = None,
+        num_layers: int = 1,
+    ) -> None:
+        super().__init__(device)
+
+        if not isinstance(module, nn.Module):
+            raise TypeError(f"module must be an nn.Module, got {type(module).__name__}")
+        if not callable(preprocess):
+            raise TypeError(
+                "preprocess must be callable (PIL image -> (3, H, W) tensor). "
+                "VisBench cannot guess a model's normalisation constants."
+            )
+
+        if num_layers < 1:
+            raise ValueError(f"num_layers must be >= 1, got {num_layers}")
+        if num_layers > 1 and layer_feature_fn is None:
+            raise ValueError(
+                f"num_layers={num_layers} promises depths VisBench cannot reach on its own. "
+                "Pass layer_feature_fn=(module, images, layers) -> [(tokens, cls, grid_hw), ...]; "
+                "an arbitrary nn.Module has no interface for tapping intermediate activations."
+            )
+        if layer_feature_fn is not None and num_layers == 1:
+            raise ValueError(
+                "layer_feature_fn was given but num_layers=1, so no multi-layer request "
+                "can ever reach it. Pass num_layers= to say how many depths it serves."
+            )
+
+        self.name = name
+        self.model = module
+        self._num_layers = num_layers
+        self._layer_feature_fn = layer_feature_fn
+        self.has_cls_token = has_cls_token
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.embed_dim = embed_dim or 0
+        self._preprocess = preprocess
+        self._feature_fn = feature_fn
+        self.weights_id = weights_id if weights_id is not None else hash_weights(module)
+
+        self._finalize()
+
+    @property
+    def num_layers(self) -> int:
+        """Depths this module exposes; 1 unless ``layer_feature_fn`` was given."""
+        return self._num_layers
+
+    def _forward_features(
+        self,
+        image: torch.Tensor,
+        layers: list[int],
+    ) -> list[LayerOutput]:
+        """Interpret the module's output as ``(patch_tokens, cls, grid_hw)``.
+
+        With ``layer_feature_fn`` supplied, that handles the whole request.
+        Otherwise exactly one layer can be asked for — the base class has
+        already rejected anything else against :attr:`num_layers`.
+
+        With ``feature_fn`` supplied, that does the work and this only
+        validates. Otherwise the module's ``forward`` output is read:
+
+        ``(B, C, H, W)``
+            A conv map. Unambiguous — the grid is the spatial shape.
+        ``(B, N, C)``
+            A token sequence. The grid is **not** recoverable from ``N`` alone,
+            so ``patch_size`` must be set, or a square grid is assumed and a
+            non-square input rejected rather than silently misaligned.
+
+        Anything else raises, because the alternative is a feature map whose
+        spatial layout is wrong in a way no shape check downstream would catch.
+        """
+        if self._layer_feature_fn is not None:
+            outputs = list(self._layer_feature_fn(self.model, image, list(layers)))
+            if len(outputs) != len(layers):
+                raise ValueError(
+                    f"layer_feature_fn returned {len(outputs)} outputs for layers={layers}. "
+                    "It must return one (tokens, cls, grid_hw) per requested index, in order."
+                )
+            result: list[LayerOutput] = []
+            for index, output in zip(layers, outputs, strict=True):
+                try:
+                    tokens, cls_token, grid_hw = output
+                except (TypeError, ValueError) as error:
+                    raise ValueError(
+                        f"layer_feature_fn's output for layer {index} is not a "
+                        f"(tokens, cls, grid_hw) triple: {output!r}"
+                    ) from error
+                self._note_embed_dim(tokens)
+                result.append((tokens, cls_token, tuple(grid_hw)))
+            return result
+
+        if self._feature_fn is not None:
+            tokens, cls_token, grid_hw = self._feature_fn(self.model, image)
+            self._note_embed_dim(tokens)
+            return [(tokens, cls_token, tuple(grid_hw))]
+
+        output = self.model(image)
+        if isinstance(output, (tuple, list)):
+            output = output[0]
+        if not isinstance(output, torch.Tensor):
+            raise TypeError(
+                f"The module returned {type(output).__name__}, not a tensor. "
+                "Pass feature_fn= to extract features yourself."
+            )
+
+        if output.ndim == 4:
+            _, _, grid_h, grid_w = output.shape
+            tokens = output.flatten(2).transpose(1, 2)
+            self._note_embed_dim(tokens)
+            return [(tokens, None, (grid_h, grid_w))]
+
+        if output.ndim == 3:
+            tokens, cls_token, grid_hw = self._tokens_to_features(output, image)
+            self._note_embed_dim(tokens)
+            return [(tokens, cls_token, grid_hw)]
+
+        raise ValueError(
+            f"The module returned a {output.ndim}D tensor {tuple(output.shape)}. "
+            "VisBench understands (B, C, H, W) conv maps and (B, N, C) token "
+            "sequences; for anything else pass feature_fn=."
+        )
+
+    def _tokens_to_features(
+        self,
+        output: torch.Tensor,
+        image: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[int, int]]:
+        """Split a token sequence into patches, CLS, and a spatial grid."""
+        cls_token = None
+        tokens = output
+        if self.has_cls_token:
+            cls_token = output[:, 0]
+            tokens = output[:, 1:]
+
+        num_tokens = tokens.shape[1]
+        _, _, height, width = image.shape
+
+        if self.patch_size is not None:
+            grid_hw = (height // self.patch_size, width // self.patch_size)
+            if grid_hw[0] * grid_hw[1] != num_tokens:
+                raise ValueError(
+                    f"patch_size={self.patch_size} implies a {grid_hw[0]}x{grid_hw[1]} grid "
+                    f"({grid_hw[0] * grid_hw[1]} tokens) but the module returned {num_tokens}. "
+                    "Check has_cls_token, or pass feature_fn=."
+                )
+            return tokens, cls_token, grid_hw
+
+        side = int(round(num_tokens**0.5))
+        if side * side != num_tokens:
+            raise ValueError(
+                f"{num_tokens} patch tokens is not a square grid, and patch_size is unset, "
+                "so the spatial layout is unknown. Set patch_size= or pass feature_fn=."
+            )
+        if height != width:
+            raise ValueError(
+                f"A square token grid was assumed, but the input is {height}x{width}. "
+                "Set patch_size= so the grid can be derived rather than guessed."
+            )
+        return tokens, cls_token, (side, side)
+
+    def _note_embed_dim(self, tokens: torch.Tensor) -> None:
+        """Fill in ``embed_dim`` from the first forward pass if it was not given.
+
+        Metadata rather than machinery — nothing depends on it — but a result
+        record saying the feature width is 0 is worse than one saying 2048.
+        """
+        if not self.embed_dim:
+            self.embed_dim = int(tokens.shape[-1])
+
+    def preprocess(self, images: Image.Image | list) -> torch.Tensor:
+        """Apply the user's transform to one image or a sequence of them."""
+        if isinstance(images, Image.Image):
+            images = [images]
+        return torch.stack([self._preprocess(img) for img in images])
+
+    def cache_key(self) -> str:
+        """``"custom/<name>/<weights_id>/<resolution>"``.
+
+        The weights id is what keeps two different checkpoints of the same
+        architecture from sharing cached features — the failure a custom
+        backbone is most exposed to, since there is no upstream ref to lean on.
+        """
+        return f"custom/{self.name}/{self.weights_id}/{self.image_size}"
+
+    def extra_repr(self) -> str:
+        return f"name={self.name!r}, weights_id={self.weights_id!r}"
+
+    def describe(self) -> dict[str, Any]:
+        """Metadata for a result record."""
+        return {
+            "backbone": self.name,
+            "backbone_key": self.cache_key(),
+            "has_cls_token": self.has_cls_token,
+            "embed_dim": self.embed_dim,
+        }