vision_unlearning.benchmarks.I_care.embeddings

DINOv2 embedding utilities for the I-CARE benchmark.

This module provides:

  • embed_forgetting_session(): embed all images from one forgetting session (entity or baseline)

  • load_dino_model(): load DINOv2 vits14 and return (model, transform, device) triple

  • embed_image_with_dino(): embed a single image using a pre-loaded DINOv2 model

Design notes:

  • Heavy GPU imports (torch, torchvision, PIL) are deferred to function call time so this module is safe to import in CPU-only environments.

  • embed_image_fn is injectable in embed_forgetting_session() for unit testing without GPU.

  • TODO: refactor embed_image_with_dino() into batched DataLoader for throughput.

  • TODO: add torch.compile() support.

  • TODO: add fp16 (half-precision) support for throughput.

  • TODO: add DataLoader parallelism (num_workers > 0).

Attributes

logger

EMBEDDING_MODEL

EMBEDDING_DIM

Functions

load_dino_model(→ Tuple[Any, Any, str])

Load DINOv2 model, transform pipeline, and device.

embed_image_with_dino(→ List[float])

Embed a single image using a pre-loaded DINOv2 model.

embed_forgetting_session(→ List[Dict[str, Any]])

Embed all images from one forgetting session (entity or baseline).

embed_forgetting_session_batched(→ List[Dict[str, Any]])

Embed all images for one forgetting session using batched GPU inference.

Module Contents

vision_unlearning.benchmarks.I_care.embeddings.logger
vision_unlearning.benchmarks.I_care.embeddings.EMBEDDING_MODEL = 'dinov2_vits14'
vision_unlearning.benchmarks.I_care.embeddings.EMBEDDING_DIM = 384
vision_unlearning.benchmarks.I_care.embeddings.load_dino_model(model_name: str = EMBEDDING_MODEL, force_device: str | None = None) → Tuple[Any, Any, str]

Load DINOv2 model, transform pipeline, and device.

Heavy imports (torch, torchvision) happen here, not at module load.

Parameters:

  • model_name – DINOv2 model variant (default: ‘dinov2_vits14’ → 384-dim CLS).

  • force_device – If set, use this device string instead of auto-detecting.

Returns:

(model, transform, device) tuple. model: DINOv2 PyTorch model in eval mode, on device. transform: torchvision.transforms pipeline (resize → crop → normalize). device: device string (‘cuda’ or ‘cpu’).

vision_unlearning.benchmarks.I_care.embeddings.embed_image_with_dino(image_path: str, model: Any, transform: Any, device: str) → List[float]

Embed a single image using a pre-loaded DINOv2 model.

Parameters:

  • image_path – Path to a PNG/JPEG image on disk.

  • model – DINOv2 model (from load_dino_model()).

  • transform – torchvision transform (from load_dino_model()).

  • device – device string (‘cuda’ or ‘cpu’).

Returns:

384-dim CLS embedding as a plain Python list of floats. TODO: refactor into batched DataLoader for throughput (currently single-image).

vision_unlearning.benchmarks.I_care.embeddings.embed_forgetting_session(dataset_folder: str, seeds: List[int], prompts: List[str], metadata_filtered: List[Dict[str, Any]], lora_state: Literal['on', 'off'], task: str, embed_image_fn: Callable[[str], List[float]] | None = None) → List[Dict[str, Any]]

Embed all images from one forgetting session (entity or baseline).

Iterates over all (seed, prompt) combinations and embeds each matching image. Images that do not exist on disk are skipped with a warning.

Parameters:

  • dataset_folder – Local directory containing the generated images.

  • seeds – List of generation seeds (e.g. [0, 1, 2, 3]).

  • prompts – Full prompt strings (e.g. “An image of Colin Powell”).

  • metadata_filtered – Metadata list used to map prompt index → entity name. metadata_filtered[i][‘name’] corresponds to prompts[i].

  • lora_state – ‘on’ for unlearned model images, ‘off’ for baseline images.

  • task – Task name, passed to get_target_preprocessed().

  • embed_image_fn – Injectable embedding function (image_path → [float]). Required — there is no default. Pass embed_image_with_dino (partially applied) or a test stub.

Returns:

[
{

‘prompted_entity’: str, # entity name (preprocessed) ‘seed’: int, ‘prompt’: str, ‘embedding’: List[float], # 384-dim CLS embedding

]

Return type:

List of records

vision_unlearning.benchmarks.I_care.embeddings.embed_forgetting_session_batched(dataset_folder: str, seeds: List[int], prompts: List[str], metadata_filtered: List[Dict[str, Any]], lora_state: Literal['on', 'off'], task: str, model: Any, transform: Any, device: str, batch_size: int = 32) → List[Dict[str, Any]]

Embed all images for one forgetting session using batched GPU inference.

More efficient than embed_forgetting_session() for large image sets. Collects all (path, metadata) pairs first, then processes in batches via a simple loop, amortising Python overhead and maximising GPU utilisation.

Parameters:

  • dataset_folder – Local directory containing the generated images.

  • seeds – List of generation seeds used.

  • prompts – Full prompt strings.

  • metadata_filtered – Metadata list: metadata_filtered[i][‘name’] → prompts[i].

  • lora_state – ‘on’ for unlearned model, ‘off’ for baseline.

  • task – Task name, passed to get_target_preprocessed().

  • model – DINOv2 model (from load_dino_model()), on device, in eval mode.

  • transform – torchvision transform pipeline (from load_dino_model()).

  • device – Torch device string (‘cuda’ or ‘cpu’).

  • batch_size – Number of images per GPU forward pass (default 32). TODO: tune based on VRAM; 32 images × 224×224 ≈ 220MB VRAM.

Returns:

Same structure as embed_forgetting_session().