`mlvbench.models`¶

Models.

This package provides a unified interface to pretrained Vision Transformers (ViTs) for extracting intermediate features.

Models are referenced by a "provider/model_name" string (e.g. "timm/vit_base_patch16_224.dino"). Use build_model to instantiate one, and list_models or list_bundles to discover what is available. The supported providers are timm, franca, webssl, and diffusers.

Model.forward_features takes a (B, 3, H, W) image batch and returns a Features container keyed by layer name. ViT features are a flat sequence of tokens — patch tokens plus optional special tokens such as cls or register tokens, whose arrangement is described by a TokenLayout. Features provides named access to these token types via .patch(), .cls(), and .register():

from mlvbench.models import build_model

model = build_model("timm/vit_base_patch16_224.dino")
features = model.forward_features(images)            # images: (B, 3, H, W)
patches = features.patch("block.11", format="BCHW")  # (B, C, H, W)
cls = features.cls("block.11")                       # (B, C)

Every model can also be built with randomly initialized weights instead of the pretrained ones, which is useful as a control condition. The architecture and all metadata stay the same; passing a seed makes the initialization reproducible:

model = build_model("timm/vit_base_patch16_224.dino", pretrained=False, seed=0)

Trained Probes can be attached to a model to read out from these features, and a model together with its probes can be round-tripped with Model.save and load_model. See mlvbench.probes for the probing API.

See the symbol reference below for full details.

`Features` ¶

Per-layer model features.

Features from Vision Transformers are a sequence of tokens that include spatial patch tokens and may include special tokens (e.g., cls or register tokens). This class is a lightweight container for model features that simplifies accessing the different token types returned by a ViT.

`layers` `property` ¶

The names of the layers held by this container.

`init(tokens, layout)` ¶

Initialize the container.

Parameters:

Name	Type	Description	Default
`tokens`	`dict[str, Tensor]`	Mapping from layer name to the full token sequence with shape `(B, N_total, C)`, arranged as described by `layout`.	required
`layout`	`TokenLayout`	The token layout shared across all layers.	required

`cls(layer)` ¶

Return the cls (class) token of a layer.

Parameters:

Name	Type	Description	Default
`layer`	`str`	The layer name.	required

Returns:

Type	Description
`Tensor`	The cls token with shape `(B, C)`.

Raises:

Type	Description
`KeyError`	If the model has no cls token.

`patch(layer, format='BNC')` ¶

Return the patch tokens of a layer.

Parameters:

Name	Type	Description	Default
`layer`	`str`	The layer name.	required
`format`	`Literal['BNC', 'BCHW']`	The output format, either `"BNC"` (default) for a flat sequence or `"BCHW"` for a spatial map.	`'BNC'`

Returns:

Type	Description
`Tensor`	The patch tokens in the requested format.

Raises:

Type	Description
`ValueError`	If `format` is unknown, or `"BCHW"` is requested but the grid size is unknown or inconsistent with the token count.

`register(layer)` ¶

Return the register tokens of a layer.

Parameters:

Name	Type	Description	Default
`layer`	`str`	The layer name.	required

Returns:

Type	Description
`Tensor`	The register tokens with shape `(B, R, C)`.

Raises:

Type	Description
`KeyError`	If the model has no register tokens.

`special(layer, name)` ¶

Return a named group of special (prefix) tokens of a layer.

Parameters:

Name	Type	Description	Default
`layer`	`str`	The layer name.	required
`name`	`str`	The name of the prefix-token group, e.g. `"cls"` or `"register"`.	required

Returns:

Type	Description
`Tensor`	The tokens of the group with shape `(B, k, C)`, where `k` is the group size.

Raises:

Type	Description
`KeyError`	If the tokens do not include a group with that name.

`subsample_patches(index)` ¶

Return a copy with the patch tokens gathered along the given index.

The same indices are applied to every layer. Prefix (special) tokens are kept unchanged, while the patch tokens are gathered, so accessors such as cls keep working.

Parameters:

Name	Type	Description	Default
`index`	`Tensor`	Patch-token indices with shape `(B, K)`, indexing into the patch tokens (i.e. excluding any prefix tokens).	required

Returns:

Type	Description
`Features`	A new Features holding the same prefix tokens and the gathered patch tokens. The grid layout is dropped (`grid_size` is None), since the subsampled patch tokens no longer form a spatial grid.

`Model` ¶

Bases: ABC, Module

Base class for models.

`compiled` `property` ¶

Whether forward_features has been compiled with torch.compile.

Detected via the marker torch.compile leaves on the wrapped callable, so it reflects the actual state regardless of how compilation was applied (e.g. by build_model).

`device` `property` ¶

The device used by the model.

`dtype` `property` ¶

The compute dtype of the model weights.

`embed_dim` `abstractmethod` `property` ¶

The number of dimensions for each token.

`input_size` `abstractmethod` `property` ¶

The native input image size the model was trained on.

`layer_names` `abstractmethod` `property` ¶

The names of the model layers.

Returns:

Type	Description
`list[str]`	A list of layer names, e.g. ["block.0", "block.1", "block.2"].

`name` `abstractmethod` `property` ¶

The name of the model.

`num_layers` `abstractmethod` `property` ¶

The number of layers in the model.

`num_parameters` `property` ¶

The number of parameters of the model.

`patch_size` `abstractmethod` `property` ¶

The patch size of the model.

`precision` `property` ¶

The requested compute precision of the model.

`pretrained` `property` ¶

Whether the model was built with pretrained weights.

If False, the architecture is identical but the weights are randomly initialized.

`seed` `property` ¶

The seed used for the random weight initialization.

`token_layout` `abstractmethod` `property` ¶

The static token layout of the model.

Describes which prefix (special) token groups the model offers, such as cls or register tokens, without running a forward pass. grid_size is None because it depends on the input resolution; the layout carried by the Features returned by forward_features fills it in.

Use this to query a model's token types ahead of time, for example:

"cls" in model.token_layout        # whether a cls token is available
model.token_layout.names           # e.g. ["cls", "register"]

`url` `abstractmethod` `property` ¶

URL pointing to more information about the model.

Ideally a model card (e.g. on Hugging Face) documenting details such as the citation and license. If no model card exists, this may instead link to the git repository or paper describing the model.

`init(pretrained=True, precision='auto', seed=0)` ¶

Initialize the model.

Parameters:

Name	Type	Description	Default
`pretrained`	`bool`	If True (default), load the pretrained weights. If False, the architecture is built with randomly initialized weights.	`True`
`precision`	`Literal['auto', 'float32', 'bfloat16']`	The compute precision of the model. `"auto"` keeps the dtype of the loaded checkpoint, `"float32"` and `"bfloat16"` cast the weights to the respective dtype.	`'auto'`
`seed`	`int`	Seed for the random weight initialization when `pretrained=False`.	`0`

`config()` ¶

Return the keyword arguments required to rebuild this model.

The returned mapping is passed to build_model to reconstruct an equivalent model, and is what Model.save persists so that a saved model can be restored with load_model.

The base implementation captures name, precision, pretrained and seed. Adapters with extra constructor options should override this and extend the returned mapping, so that those options survive a save/load round-trip:

def config(self) -> dict[str, Any]:
    return {**super().config(), "timestep": self._timestep}

`forward_features(images, layers=None)` `abstractmethod` ¶

Extract intermediate features for the images.

Parameters:

Name	Type	Description	Default
`images`	`Tensor`	Input images with shape `(B, 3, H, W)` and dtype `uint8` or `float32`.	required
`layers`	`list[str] \| None`	The list of layer names to extract features from. If None, features from all layers are extracted.	`None`

Returns:

Type	Description
`Features`	Per-layer features extracted from the model, always with dtype `float32` regardless of the model's compute precision. The `Features` container provides named access to the different token types, for example: `features.patch("block.11") # (B, H*W, C) features.patch("block.11", format="BCHW") # (B, C, H, W) features.cls("block.11") # (B, C)`

`forward_probes(images)` ¶

Extract features and apply the attached probes.

Parameters:

Name	Type	Description	Default
`images`	`Tensor`	Input images with shape `(B, 3, H, W)` and dtype `uint8` or `float32`.	required

Returns:

Type	Description
`list[Tensor]`	A list of predictions in the order of `self.probes`.

`forward_probes_from_features(features)` ¶

Apply each attached probe to the given features.

Parameters:

Name	Type	Description	Default
`features`	`Features`	Per-layer features as returned by `forward_features`.	required

Returns:

Type	Description
`list[Tensor]`	A list of predictions in the order of `self.probes`.

`info()` ¶

Return information about the model.

`save(path, include_weights=False)` ¶

Save this model's state so it can be restored with load_model.

Stores the build configuration (so the pretrained backbone is rebuilt by name via config) and all attached probes. Backbone weights are not stored by default — they are reconstructed from the pretrained checkpoint, which keeps the file small.

Each probe stores its own name, so a model may hold probes of different types. Custom probe types (those not built in) must be supplied via the probe_types argument of load_model to rebuild the probes.

Parameters:

Name	Type	Description	Default
`path`	`Path`	Destination file path.	required
`include_weights`	`bool`	If True, additionally store the backbone weights (its `state_dict`). Use this when the weights have been modified after loading; otherwise the pretrained checkpoint is sufficient.	`False`

`TokenLayout` `dataclass` ¶

Describes the token layout in Vision Transformers features.

Vision Transformers process a flat sequence of tokens that includes the spatial patch tokens, and optionally special tokens (cls, register, ...). The TokenLayout describes how the tokens are arranged in the sequences.

For example, if a model uses 1 cls token, 4 register tokens followed by patch tokens on a 10x10 grid, we exepct 105 tokens in total. This arrangement is described by the following layout:

TokenLayout(
    prefix=(("cls", 1), ("register", 4)),
    grid_size=(10, 10),
)

Attributes:

Name	Type	Description
`prefix`	`tuple[tuple[str, int], ...]`	Ordered prefix-token groups as `(name, count)` pairs, e.g. `(("cls", 1), ("register", 4))`. Empty if the model has no special tokens.
`grid_size`	`tuple[int, int] \| None`	The `(H, W)` spatial grid size of the patch tokens, used to reshape patch tokens to `BCHW`. `None` if unknown.

`names` `property` ¶

The names of the prefix-token groups, in token order.

Returns:

Type	Description
`list[str]`	The group names, e.g. `["cls", "register"]`. Empty if the model has no special tokens.

`num_prefix_tokens` `property` ¶

The total number of prefix (special) tokens across all groups.

`contains(name)` ¶

Whether the layout declares a prefix-token group with the given name.

Enables membership tests such as "cls" in model.token_layout.

Parameters:

Name	Type	Description	Default
`name`	`str`	The name of a prefix-token group, e.g. `"cls"` or `"register"`.	required

Returns:

Type	Description
`bool`	`True` if a group with this name is declared, `False` otherwise.

`span(name)` ¶

Return the [start, end) index range of a prefix group.

Parameters:

Name	Type	Description	Default
`name`	`str`	The name of the prefix-token group (e.g., `"cls"` or `"register"`),	required

Returns:

Type	Description
`tuple[int, int]`	The start (inclusive) and end (exclusive) token indices of the group.

Raises:

Type	Description
`KeyError`	If no prefix group with this name is declared.

`build_model(name, pretrained=True, precision='auto', device='auto', compile=False, seed=0, **kwargs)` ¶

Build the model with the given name.

Parameters:

Name	Type	Description	Default
`name`	`str`	Model name in "provider/model_name" format.	required
`pretrained`	`bool`	If True (default), load the pretrained weights. If False, the same architecture is built with randomly initialized weights.	`True`
`precision`	`Literal['auto', 'float32', 'bfloat16']`	The compute precision of the model. `"auto"` keeps the checkpoint dtype, `"float32"` and `"bfloat16"` cast the weights accordingly.	`'auto'`
`device`	`device \| str`	The device to load the model onto. `"auto"` selects CUDA if available and falls back to the CPU.	`'auto'`
`compile`	`bool`	If True, compile the model's `forward_features` with `torch.compile`. Only the base model's feature extraction is compiled; attached probes are invoked separately and remain eager. The returned object stays a plain `Model` instance.	`False`
`seed`	`int`	Seed for the random weight initialization when `pretrained=False`.	`0`
`**kwargs`	`Any`	Additional model-specific options forwarded to the model constructor. Passing an option that the selected model does not support raises a `TypeError`.	`{}`

Returns:

Type	Description
`Model`	The model instance.

`list_bundles()` ¶

Return all available model bundles.

`list_models(patterns=None, bundle=None)` ¶

List available models.

Parameters:

Name	Type	Description	Default
`patterns`	`list[str] \| None`	Patterns to filter model names, supporting glob-like wildcards. If provided, only models matching any of the patterns will be returned.	`None`
`bundle`	`str \| None`	Name of the model bundle to list models from. If provided, only models in the bundle will be returned.	`None`

Returns:

Type	Description
`list[str]`	List of model names in "provider/model_name" format.

`load_model(path, device='auto', compile=False, probe_types=None)` ¶

Load a model with probes from a checkpoint file.

Parameters:

Name	Type	Description	Default
`path`	`Path \| str`	Path to a checkpoint written by `Model.save`.	required
`device`	`device \| str`	The device to load the model and probes onto. "auto" selects CUDA if available and falls back to the CPU.	`'auto'`
`compile`	`bool`	If True, compile the base model's `forward_features` with `torch.compile`. The attached probes remain eager.	`False`
`probe_types`	`dict[str, type[Probe]] \| None`	Optional mapping of custom probe type names to `Probe` subclasses. Required if the checkpoint contains probes of a custom type.	`None`

Returns:

Type	Description
`Model`	The reconstructed model with its probes attached.

Raises:

Type	Description
`ValueError`	If the checkpoint's format version is not supported by this mlvbench, or if a stored probe name is neither a built-in type nor present in `probe_types`.

mlvbench.models¶

Features ¶

layers property ¶

__init__(tokens, layout) ¶

cls(layer) ¶

patch(layer, format='BNC') ¶

register(layer) ¶

special(layer, name) ¶

subsample_patches(index) ¶

Model ¶

compiled property ¶

device property ¶

dtype property ¶

embed_dim abstractmethod property ¶

input_size abstractmethod property ¶

layer_names abstractmethod property ¶

name abstractmethod property ¶

num_layers abstractmethod property ¶

num_parameters property ¶

patch_size abstractmethod property ¶

precision property ¶

pretrained property ¶

seed property ¶

token_layout abstractmethod property ¶

url abstractmethod property ¶

__init__(pretrained=True, precision='auto', seed=0) ¶

config() ¶

forward_features(images, layers=None) abstractmethod ¶

forward_probes(images) ¶

forward_probes_from_features(features) ¶

info() ¶

save(path, include_weights=False) ¶

TokenLayout dataclass ¶

names property ¶

num_prefix_tokens property ¶

__contains__(name) ¶

span(name) ¶

build_model(name, pretrained=True, precision='auto', device='auto', compile=False, seed=0, **kwargs) ¶

list_bundles() ¶

list_models(patterns=None, bundle=None) ¶

load_model(path, device='auto', compile=False, probe_types=None) ¶

`mlvbench.models`¶

`Features` ¶

`layers` `property` ¶

`init(tokens, layout)` ¶

`cls(layer)` ¶

`patch(layer, format='BNC')` ¶

`register(layer)` ¶

`special(layer, name)` ¶

`subsample_patches(index)` ¶

`Model` ¶

`compiled` `property` ¶

`device` `property` ¶

`dtype` `property` ¶

`embed_dim` `abstractmethod` `property` ¶

`input_size` `abstractmethod` `property` ¶

`layer_names` `abstractmethod` `property` ¶

`name` `abstractmethod` `property` ¶

`num_layers` `abstractmethod` `property` ¶

`num_parameters` `property` ¶

`patch_size` `abstractmethod` `property` ¶

`precision` `property` ¶

`pretrained` `property` ¶

`seed` `property` ¶

`token_layout` `abstractmethod` `property` ¶

`url` `abstractmethod` `property` ¶

`init(pretrained=True, precision='auto', seed=0)` ¶

`config()` ¶

`forward_features(images, layers=None)` `abstractmethod` ¶

`forward_probes(images)` ¶

`forward_probes_from_features(features)` ¶

`info()` ¶

`save(path, include_weights=False)` ¶

`TokenLayout` `dataclass` ¶

`names` `property` ¶

`num_prefix_tokens` `property` ¶

`contains(name)` ¶

`span(name)` ¶

`build_model(name, pretrained=True, precision='auto', device='auto', compile=False, seed=0, **kwargs)` ¶

`list_bundles()` ¶

`list_models(patterns=None, bundle=None)` ¶

`load_model(path, device='auto', compile=False, probe_types=None)` ¶