MS COCO Tutorial

This tutorial links images and captions in MS COCO. Each image is one data view, and each of its five human-written captions is a separate language view. DLVPM learns latent variables shared between the image and caption modalities, then the learned representation is evaluated with an image-text retrieval task.

The full runnable script is deep_lvpm.tutorial.tutorial_coco.

Run it with:

python -m deep_lvpm.tutorial.tutorial_coco

Prerequisites

This is the heaviest tutorial. It uses torchvision image models, Hugging Face transformers, fiftyone for COCO access, and the standard scientific Python stack. A GPU or Apple Silicon machine is strongly recommended.

1. Configure the benchmark

The tutorial keeps the important settings near the top of the script. The main modelling choice is to treat each sample as six views: one image and five caption views.

import os
import json
import random
import zipfile
import urllib.request
from collections import defaultdict

os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("USE_TORCH", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchvision.models import EfficientNet_B0_Weights, efficientnet_b0
from transformers import AutoModel, AutoTokenizer
import fiftyone as fo
import fiftyone.zoo as foz

from deep_lvpm.model import StructuralModel
from deep_lvpm.multi_model import CLIP, VICReg, LeJEPA

NUM_CAPTION_VIEWS = 5
IMG_SIZE = 224
MAX_TOKENS = 32
TEXT_MODEL_NAME = "distilbert-base-uncased"
TEXT_DROPOUT = 0.10

NDIMS = 512
BATCH_SIZE = 512
BENCHMARK_EPOCHS = 30

BENCHMARK_TRAIN_SAMPLES = 20000
BENCHMARK_VAL_SAMPLES = 5000
BENCHMARK_SAMPLES = 2048

NUM_WORKERS = 0
RUN_BASELINES = False
RETRIEVAL_KS = (1, 5, 10)
LEARNING_RATE_START = 1e-5
LEARNING_RATE_END = 1e-4
LEARNING_RATE_WARMUP_EPOCHS = 5
TEST_FRACTION = 0.10
SEED = 51

2. Define the image-caption path model

The path matrix is star-shaped. The image view is connected to every caption view. Captions are not directly connected to one another, because the shared structure should pass through the image.

Path = np.array(
    [
        [0, 1, 1, 1, 1, 1],
        [1, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0],
    ],
    dtype="float32",
)

3. Load COCO images and captions

The image splits come from the FiftyOne COCO zoo integration. The caption JSON files are resolved locally when available and downloaded from the official COCO annotation archive when needed.

train_view = foz.load_zoo_dataset(
    "coco-2017",
    split="train",
    label_types=[],
    shuffle=True,
    seed=SEED,
    dataset_name="dlvpm-coco2017-train",
    include_id=True,
)

val_view = foz.load_zoo_dataset(
    "coco-2017",
    split="validation",
    label_types=[],
    shuffle=True,
    seed=SEED,
    dataset_name="dlvpm-coco2017-val",
    include_id=True,
)

The helper functions in the script do three jobs: find the annotation JSON, load captions into a dictionary keyed by image id, and turn each image into one filepath plus five captions.

def resolve_coco_caption_annotations(split: str) -> str:
    # Locate captions_train2017.json or captions_val2017.json.
    # If the file is not present locally, download and extract the official
    # COCO annotations archive.
    ...

def load_coco_captions(annotation_path: str) -> dict[int, list[str]]:
    with open(annotation_path, "r", encoding="utf-8") as f:
        payload = json.load(f)

    captions_by_image_id = defaultdict(list)
    for annotation in payload["annotations"]:
        captions_by_image_id[int(annotation["image_id"])].append(annotation["caption"])
    return captions_by_image_id

def coco_view_to_examples(view, captions_by_image_id):
    image_paths = []
    caption_sets = []
    for sample in view:
        image_id = extract_coco_image_id(sample)
        captions = captions_by_image_id.get(image_id, [])
        if len(captions) < NUM_CAPTION_VIEWS:
            continue
        image_paths.append(sample.filepath)
        caption_sets.append(captions[:NUM_CAPTION_VIEWS])
    return image_paths, caption_sets

4. Tokenize caption views

The caption encoder is DistilBERT, so each caption view is tokenized into input_ids and attention_mask arrays. Each sample has five captions, and each caption is treated as its own DLVPM view.

tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME)

def tokenize_caption_sets(caption_sets: list[list[str]]) -> tuple[np.ndarray, np.ndarray]:
    flat_captions = [caption for caption_set in caption_sets for caption in caption_set]
    tokenized = tokenizer(
        flat_captions,
        padding="max_length",
        truncation=True,
        max_length=MAX_TOKENS,
        return_tensors="np",
    )
    input_ids = tokenized["input_ids"].reshape(len(caption_sets), NUM_CAPTION_VIEWS, MAX_TOKENS)
    attention_mask = tokenized["attention_mask"].reshape(len(caption_sets), NUM_CAPTION_VIEWS, MAX_TOKENS)
    return input_ids.astype("int64"), attention_mask.astype("int64")

5. Define image and text measurement models

The image measurement model uses EfficientNet-B0 and replaces the classifier with a projection to NDIMS. The caption measurement model uses DistilBERT and projects the pooled text representation to the same width.

class TextEncoderModule(nn.Module):
    def __init__(self, model_name: str) -> None:
        super().__init__()
        self.text_model = AutoModel.from_pretrained(model_name)
        hidden_size = int(self.text_model.config.hidden_size)
        self.dropout = nn.Dropout(TEXT_DROPOUT)
        self.projection = nn.Linear(hidden_size, NDIMS)
        self.n_inputs = 2

    def forward(self, inputs):
        input_ids, attention_mask = inputs
        outputs = self.text_model(
            input_ids=input_ids.long(),
            attention_mask=attention_mask.long(),
        )
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        cls_embedding = self.dropout(cls_embedding)
        return self.projection(cls_embedding)

class ImageEncoderModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        weights = EfficientNet_B0_Weights.DEFAULT
        self.transforms = weights.transforms()
        self.backbone = efficientnet_b0(weights=weights)
        in_features = int(self.backbone.classifier[1].in_features)
        self.backbone.classifier = nn.Sequential(nn.Dropout(p=0.2), nn.Linear(in_features, NDIMS))
        self.n_inputs = 1

    def forward(self, images):
        return self.backbone(images)

def build_model_list(NDIMS) -> list[nn.Module]:
    image_model = ImageEncoderModule()
    text_models = [TextEncoderModule(TEXT_MODEL_NAME) for _ in range(NUM_CAPTION_VIEWS)]
    return [image_model] + text_models

6. Build multiview dataloaders

Each batch contains one image tensor plus five pairs of caption tensors. The n_inputs attribute on the text encoder tells StructuralModel that each caption view consumes two tensors.

class CocoRetrievalDataset(Dataset):
    def __init__(self, image_paths, caption_sets, training=True):
        self.image_paths = list(image_paths)
        self.caption_sets = list(caption_sets)
        self.training = bool(training)
        self.input_ids, self.attention_mask = tokenize_caption_sets(self.caption_sets)

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, index):
        image = Image.open(self.image_paths[index]).convert("RGB")
        image = image.resize((IMG_SIZE, IMG_SIZE), RESAMPLE_BICUBIC)
        image = np.asarray(image).astype("float32") / 255.0
        image = np.transpose(image, (2, 0, 1))

        caption_inputs = []
        for caption_index in range(NUM_CAPTION_VIEWS):
            caption_inputs.append(self.input_ids[index, caption_index])
            caption_inputs.append(self.attention_mask[index, caption_index])
        return tuple([image] + caption_inputs)

def collate_multiview(batch):
    columns = list(zip(*batch))
    return tuple(torch.as_tensor(np.stack(column)) for column in columns)

def make_multiview_loader(image_paths, caption_sets, training=True):
    dataset = CocoRetrievalDataset(image_paths, caption_sets, training=training)
    return DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        shuffle=bool(training),
        num_workers=NUM_WORKERS,
        collate_fn=collate_multiview,
        drop_last=bool(training),
    )

7. Train DLVPM

Before training, the script builds smaller benchmark splits from the full COCO views. This keeps the tutorial practical while preserving a held-out test subset for retrieval evaluation.

benchmark_train_n = min(BENCHMARK_TRAIN_SAMPLES, len(train_paths))
benchmark_val_n = min(BENCHMARK_VAL_SAMPLES, len(val_paths))
benchmark_test_n = min(BENCHMARK_SAMPLES, len(test_paths))

benchmark_train_ds = make_multiview_loader(
    train_paths[:benchmark_train_n],
    train_caption_sets[:benchmark_train_n],
    training=True,
)
benchmark_val_ds = make_multiview_loader(
    val_paths[:benchmark_val_n],
    val_caption_sets[:benchmark_val_n],
    training=False,
)
benchmark_test_ds = make_multiview_loader(
    test_paths[:benchmark_test_n],
    test_caption_sets[:benchmark_test_n],
    training=False,
)

The benchmark model uses ZCA orthogonalization and ordered DLVs. The ordered dimension cutoff keeps the smallest number of ordered dimensions whose cumulative association mass reaches the cutoff.

def make_optimizer_list(model: nn.Module) -> list[torch.optim.Optimizer]:
    return [
        torch.optim.Adam(view_model.parameters(), lr=LEARNING_RATE_START)
        for view_model in model.model_list
    ]

def make_scheduler_list(optimizers: list[torch.optim.Optimizer]) -> list[torch.optim.lr_scheduler.LambdaLR]:
    def schedule(epoch: int) -> float:
        if epoch < LEARNING_RATE_WARMUP_EPOCHS:
            return 1.0
        return LEARNING_RATE_END / LEARNING_RATE_START

    return [torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=schedule) for optimizer in optimizers]

dlvpm_benchmark_models = build_model_list(NDIMS)
dlvpm_benchmark = StructuralModel(
    Path=Path,
    model_list=dlvpm_benchmark_models,
    regularizer_list=[None for _ in dlvpm_benchmark_models],
    tot_num=benchmark_train_n,
    ndims=NDIMS,
    orthogonalization="zca",
    diag_offset=1e-6,
    train_DLV=True,
    momentum=0.95,
    order=True,
    order_association_cutoff=0.99,
)

dlvpm_optimizers = make_optimizer_list(dlvpm_benchmark)
dlvpm_schedulers = make_scheduler_list(dlvpm_optimizers)
dlvpm_benchmark.compile(dlvpm_optimizers)
dlvpm_benchmark.fit(
    benchmark_train_ds,
    validation_data=benchmark_val_ds,
    epochs=BENCHMARK_EPOCHS,
    verbose=True,
    schedulers=dlvpm_schedulers,
)

8. Evaluate image-text retrieval

After training, the script collects image embeddings and aggregated caption embeddings, normalizes them, and computes bidirectional retrieval ranks. The reported metrics include recall at k and median rank.

def l2_normalize(matrix: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    denom = np.linalg.norm(matrix, axis=1, keepdims=True)
    return matrix / np.maximum(denom, eps)

def aggregate_caption_groups(text_embeddings: np.ndarray) -> np.ndarray:
    return text_embeddings.reshape(-1, NUM_CAPTION_VIEWS, text_embeddings.shape[-1]).mean(axis=1)

def retrieval_metrics_from_ranks(ranks: np.ndarray) -> dict[str, float]:
    metrics = {f"recall@{k}": float(np.mean(ranks <= k)) for k in RETRIEVAL_KS}
    metrics["median_rank"] = float(np.median(ranks))
    metrics["mean_rank"] = float(np.mean(ranks))
    return metrics

dlvpm_img, dlvpm_txt = collect_image_text_embeddings(
    dlvpm_benchmark,
    benchmark_test_ds,
    max_samples=benchmark_test_n,
)
benchmark_results["DLVPM"], benchmark_rank_results["DLVPM"] = evaluate_retrieval_result(
    dlvpm_img,
    dlvpm_txt,
)

9. Optional baselines

If RUN_BASELINES is set to True, the same encoders and data are used to train CLIP, VICReg, and LeJEPA baselines from deep_lvpm.multi_model. This is intended for benchmarking rather than for the minimum DLVPM workflow.

if RUN_BASELINES:
    for method_name, method_class in [
        ("CLIP", CLIP),
        ("VICReg", VICReg),
        ("LeJEPA", LeJEPA),
    ]:
        baseline_models = build_model_list(NDIMS)
        baseline_model = method_class(
            baseline_models,
            [None for _ in baseline_models],
            NDIMS,
        )
        baseline_optimizers = make_optimizer_list(baseline_model)
        baseline_model.compile(baseline_optimizers)
        baseline_model.fit(
            benchmark_train_ds,
            validation_data=benchmark_val_ds,
            epochs=BENCHMARK_EPOCHS,
            verbose=True,
        )

Summary

The COCO tutorial is a full multimodal example: it builds six aligned data views, uses pretrained image and text encoders as PyTorch measurement models, trains DLVPM with an image-caption path matrix, and evaluates whether the learned latent space supports retrieval.