embedata

Unified API for loading 29 vision datasets and extracting embeddings from any PyTorch model. Handles torchvision, HuggingFace, and custom ImageFolder datasets through a single interface. Ships with 12 built-in embedding models (CLIP, DINOv2, DINOv3) and supports async extraction on secondary GPUs.

Install

pip install embedata              # core (torchvision datasets)
pip install embedata[hf]          # + HuggingFace datasets (ImageNet, CUB, RESISC45)

From source:

git clone https://github.com/ayghri/embedata.git && cd embedata
pip install -e ".[hf]"

Quick start

from embedata import list_datasets, get_dataset, get_datasets, get_dataloaders

print(list_datasets())  # all 29 datasets

train_ds = get_dataset("cifar10", split="train", root_dir="./data")
train_ds, val_ds = get_datasets("cifar10", root_dir="./data")
train_loader, val_loader = get_dataloaders("cifar10", batch_size=128, root_dir="./data")

Most datasets auto-download on first use. Datasets that require manual setup (Kaggle downloads, frame extraction, etc.) document their steps via get_spec("dataset_name").notes.

Extracting embeddings

Pass any PyTorch model that maps images to feature vectors. Embeddings are saved as .npy files:

import torch
from embedata import get_dataloaders, extract

model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitg14")
model.eval().to("cuda:0")

train_loader, val_loader = get_dataloaders("cifar10", batch_size=256, root_dir="./data")

extract(train_loader, model, device=torch.device("cuda:0"),
        save_dir="./representations/cifar10/dinov2", suffix="train")
extract(val_loader, model, device=torch.device("cuda:0"),
        save_dir="./representations/cifar10/dinov2", suffix="val")

Load them back as a PyTorch Dataset:

from embedata import load_embeddings

ds = load_embeddings("cifar10", "dinov2", repr_dir="./representations", split="train")
feat, label = ds[0]

For on-the-fly extraction without disk I/O, EmbeddingDataLoader runs a model on a secondary device in a background thread:

from embedata import EmbeddingDataLoader, get_dataset

dataset = get_dataset("cifar10", split="train", root_dir="./data")

loader = EmbeddingDataLoader(dataset, model, device="cuda:1", batch_size=256, prefetch=2)

for embeddings, labels in loader:
    ...

Built-in models

The package includes 12 pre-registered models accessible via load_model. Eight CLIP variants (clipRN50, clipRN101, clipRN50x4, clipRN50x16, clipRN50x64, clipvitB32, clipvitB16, clipvitL14), DINOv2 ViT-g/14 (dinov2), and three DINOv3 variants (dinov3s, dinov3b, dinov3l).

from embedata import list_models, load_model

model, preprocess = load_model("clipvitL14", device="cuda:0", models_dir="./models")

The returned preprocess is the model's image transform -- pass it to get_dataloaders or get_dataset. For DINOv2, preprocess is None; use get_default_transforms() instead.

CLIP models require pip install git+https://github.com/openai/CLIP.git. DINOv3 models require pip install transformers.

Custom datasets

Register your own dataset loader with the @register decorator. The function receives a torchvision transform and a data path, and returns (train_dataset, val_dataset):

from embedata import register

@register("my_dataset", notes="Setup instructions shown by get_spec()")
def _my_dataset(transform, data_path):
    train_ds = ...
    val_ds = ...
    return train_ds, val_ds

HuggingFace streaming

Three datasets (ImageNet, CUB, RESISC45) load via HuggingFace and support streaming mode (requires pip install embedata[hf]). ImageNet-1k is gated and requires huggingface-cli login.

train_ds, val_ds = get_datasets("imagenet", streaming=True)

Dataset preparation

Some datasets require manual download and preparation before use. Python prepare scripts are bundled with the package:

python -m embedata.prepare.birdsnap     --root_dir ROOT_DIR
python -m embedata.prepare.fer2013      --root_dir ROOT_DIR
python -m embedata.prepare.ucf101       --root_dir ROOT_DIR [--download]
python -m embedata.prepare.hatefulmemes --root_dir ROOT_DIR

Shell scripts for cars, eurosat, and sun397 are included under embedata/prepare/. All scripts expect raw data under ROOT_DIR/datasets/{dataset_name}/ and write prepared splits to the same location.

Available datasets

Counts reflect splits as loaded by embedata (some merge train+val for training). "variable" means images have different native resolutions -- all are resized by the transform (default 224x224). Datasets marked "auto-download" are fetched on first use.

License

MIT