Implementation:NVIDIA NeMo Curator ImageEmbeddingStage
| Metadata | |
|---|---|
| Knowledge Sources | Paper: Learning Transferable Visual Models |
| Domains | Data_Curation, Image_Processing, Representation_Learning |
| Last Updated | 2026-02-14 |
Overview
ImageEmbeddingStage is a processing stage that computes dense vector representations of images using the CLIP ViT-L/14 model, populating embedding arrays in ImageBatch objects for downstream filtering and deduplication.
Description
ImageEmbeddingStage is a dataclass-based processing stage that implements the ProcessingStage[ImageBatch, ImageBatch] interface. It takes ImageBatch objects containing decoded image data as NumPy arrays and computes CLIP embeddings for each image using GPU-accelerated inference. The stage loads the CLIP ViT-L/14 model from a configurable model directory and processes images in configurable inference batch sizes. The computed embedding vectors are stored in the embedding field of each ImageObject within the batch. Optionally, the raw image data can be removed after embedding computation to reduce memory usage in downstream stages that only require the embedding vectors.
Usage
Use ImageEmbeddingStage after ImageReaderStage and before any filtering or deduplication stages. Configure model_dir to point to the directory containing CLIP model weights. Adjust model_inference_batch_size based on available GPU memory. Set remove_image_data=True when downstream stages do not need raw pixel data.
Code Reference
Source Location
nemo_curator/stages/image/embedders/clip_embedder.py, lines 28-117.
Signature
@dataclass
class ImageEmbeddingStage(ProcessingStage[ImageBatch, ImageBatch]):
model_dir: str = None
num_gpus_per_worker: float = 0.25
model_inference_batch_size: int = 32
verbose: bool = False
remove_image_data: bool = False
name: str = "image_embedding"
Import
from nemo_curator.stages.image.embedders.clip_embedder import ImageEmbeddingStage
I/O Contract
| Direction | Type | Description |
|---|---|---|
| Input | ImageBatch
|
An ImageBatch containing ImageObject instances with image_data as NumPy arrays in [H, W, C] RGB format.
|
| Output | ImageBatch
|
An ImageBatch with embedding NumPy arrays populated for each ImageObject. If remove_image_data=True, the image_data field is cleared.
|
Usage Examples
from nemo_curator.stages.image.embedders.clip_embedder import ImageEmbeddingStage
# Create embedding stage with default settings
embedder = ImageEmbeddingStage(
model_dir="/path/to/clip/weights",
model_inference_batch_size=32,
num_gpus_per_worker=0.25,
)
# Create embedding stage that removes image data after embedding
embedder_lean = ImageEmbeddingStage(
model_dir="/path/to/clip/weights",
model_inference_batch_size=64,
remove_image_data=True,
verbose=True,
name="clip_embedder",
)