CLIP Training, Evaluation, Inference, and Export

The following sections cover the experiment specification parameters, training and evaluation commands, inference, export, and TRT deployment for CLIP.

For an overview of supported models, data formats, and end-to-end workflows, refer to CLIP Introduction.

Creating an Experiment Specification File

model

model:
  type: siglip2-so400m-patch16-256
  adaptor_name: null
  freeze_vision_encoder: false
  freeze_text_encoder: false
  canonicalize_text: false

Field	Data type	Description	Default value	Valid options
model.type	string	Backbone architecture. Refer to CLIP Introduction for all valid values.	siglip2-so400m-patch16-256
model.adaptor_name	string	Text adaptor for Radio-CLIP models. Set to siglip or clip. Required when using a Radio-CLIP backbone.	null	siglip, clip
model.freeze_vision_encoder	bool	Freeze vision encoder weights during training. When peft.enabled is true, peft.vision.mode overrides this flag. Refer to PEFT Configuration.	false	true, false
model.freeze_text_encoder	bool	Freeze text encoder weights during training. When peft.enabled is true, peft.text.mode overrides this flag. Refer to PEFT Configuration.	false	true, false
model.canonicalize_text	bool	Lowercase and remove punctuation from captions before tokenization. Enable this only if the pretrained model was trained with text canonicalization.	false	true, false

train

train:
  num_epochs: 100
  num_gpus: 1
  num_nodes: 1
  checkpoint_interval: 10
  loss_type: siglip
  siglip_loss_dist_impl: gather
  siglip_loss_mask_mode: none
  compatible_positive_weight: 1.0
  compatible_positive_normalization: per_query
  triplet_loss_weight: 0.0
  triplet_margin: 0.2
  precision: fp16
  distributed_strategy: ddp
  grad_checkpointing: false
  grad_clip_norm: null
  pretrained_model_path: null
  resume_training_checkpoint_path: null
  validation_interval: 1
  val_check_interval: null

Field	Data type	Description	Default value	Valid options
train.num_epochs	int	Total number of training epochs.	100
train.num_gpus	int	Number of GPUs per node.	1
train.num_nodes	int	Number of nodes for distributed training.	1
train.checkpoint_interval	int	Save a checkpoint every N epochs.	10
train.loss_type	string	Contrastive loss formulation. Use siglip for sigmoid-based loss — the recommended choice for SigLIP2 and Radio-CLIP — or clip for softmax-based loss.	siglip	siglip, clip
train.siglip_loss_dist_impl	categorical	Distributed implementation of the negative exchange in the Sigmoid Loss for Language Image Pre-Training (SigLIP). TAO uses this field only when train.loss_type is siglip. Use local to disable negative exchange across processes. The metadata-aware masking modes support only local and gather. Refer to Metadata-Aware SigLIP Loss Masking.	gather	bidir, shift, reduce, gather, local
train.siglip_loss_mask_mode	categorical	Metadata-based masking mode for the SigLIP loss. The value none keeps the standard SigLIP loss. With train.loss_type: clip, TAO ignores this field. Refer to Metadata-Aware SigLIP Loss Masking for the modes and their requirements.	none	none, attribute_match_ignore, attribute_plus_accessory_match_ignore, attribute_match_positive, attribute_plus_accessory_match_positive
train.compatible_positive_weight	float	Weight of the promoted pairs, which are the compatible off-diagonal pairs that TAO treats as additional positive pairs. Set to 0.0 to give promoted pairs no loss weight. This field has an effect only with the attribute_match_positive and attribute_plus_accessory_match_positive modes. Refer to Metadata-Aware SigLIP Loss Masking.	1.0	min 0.0
train.compatible_positive_normalization	categorical	Apply compatible_positive_weight to every promoted pair with per_pair, or divide it among the promoted images of each text query with per_query. Refer to Metadata-Aware SigLIP Loss Masking for how TAO counts the promoted images.	per_query	per_pair, per_query
train.triplet_loss_weight	float	Weight for the auxiliary batch-hard image-text triplet loss. Set to 0.0 to disable. The positive modes require 0.0. Refer to Metadata-Aware SigLIP Loss Masking.	0.0	min 0.0
train.triplet_margin	float	Margin for the auxiliary batch-hard image-text triplet loss.	0.2	min 0.0
train.precision	string	Training precision.	fp16	fp16, bf16, fp32
train.distributed_strategy	string	Distributed training strategy. Use fsdp for very large models that exceed single-GPU memory.	ddp	ddp, fsdp
train.grad_checkpointing	bool	Enable gradient checkpointing to reduce GPU memory at the cost of additional compute.	false	true, false
train.grad_clip_norm	float	Maximum gradient norm for clipping. Set to null to disable.	null
train.pretrained_model_path	string	Path to a TAO checkpoint to use as the starting point for fine-tuning. When set to null, TAO loads pretrained weights from HuggingFace or torch.hub.	null
train.resume_training_checkpoint_path	string	Path to a TAO checkpoint from which to resume an interrupted training run.	null
train.validation_interval	int	Run validation every N epochs.	1
train.val_check_interval	int	Run validation every N steps. When set, this takes precedence over validation_interval.	null

Metadata-Aware SigLIP Loss Masking

The standard SigLIP loss compares every image in a batch with every caption in the batch. A pair of an image and its own caption is the paired positive. Every other pair is an off-diagonal pair, which the loss treats as a negative. The train.siglip_loss_mask_mode field changes how TAO treats the off-diagonal pairs whose text query matches the image. Each caption in the batch acts as a text query. TAO reads the metadata of the images and queries from the file that dataset.train.datasets[].train_pairs_file names. The value none keeps the standard loss. The other four values enable metadata-aware masking. Off-diagonal pairs that do not match remain negatives, and TAO always keeps the paired positive.

An off-diagonal pair is compatible when the text query of the pair matches the image. TAO decides whether a text query matches an image by these rules:

• Every attribute value that the query specifies must equal the attribute value of the image.

• An attribute that the query does not specify matches any image value.

• An attribute with an unknown value for the image never satisfies a value that the query specifies.

An unknown attribute value is a negative integer. TAO converts a label that the attribute vocabulary names not visible to a negative value only when an attribute_vocab.json file sits next to the pairs file. Without that file, TAO logs a warning and treats the label as a literal ID.

attribute_plus_accessory_match_ignore and attribute_plus_accessory_match_positive add one rule. Every accessory ID of the query must appear in the accessory IDs of the image. A query that has no accessory IDs matches every image on this rule.

A query has no metadata constraint when it specifies no attributes. For the two accessory modes, it must also specify no accessory IDs. Such a query matches every image. For the ignore modes, TAO therefore drops all off-diagonal pairs of that query from the loss. For the positive modes, TAO promotes no pair for that query and also drops its off-diagonal pairs from the loss, so only the paired positive contributes.

The modes whose names end in _ignore are the ignore modes, and the modes whose names end in _positive are the positive modes.

Mode	Treatment of compatible pairs
attribute_match_ignore	TAO drops the compatible pairs from the loss. A pair is compatible according to the attribute rules only.
attribute_plus_accessory_match_ignore	TAO drops the compatible pairs from the loss. A pair is compatible according to the attribute and accessory ID rules.
attribute_match_positive	TAO promotes the compatible pairs of each query that specifies at least one attribute, which means that TAO treats them as additional positive pairs. A pair is compatible according to the attribute rules only.
attribute_plus_accessory_match_positive	TAO promotes the compatible pairs of each query that specifies at least one attribute or accessory ID, in the same way. A pair is compatible according to the attribute and accessory ID rules.

The paired positive and the remaining negatives keep a weight of 1. The positive modes weight the promoted pairs according to train.compatible_positive_normalization:

• per_pair: Each promoted pair receives the weight train.compatible_positive_weight.

• per_query: TAO divides the weight among the promoted images of each text query. With train.siglip_loss_dist_impl: gather and more than one process, TAO counts the promoted images across all processes. With train.siglip_loss_dist_impl: local, TAO counts only the local batch.

For the positive modes, TAO logs the train/compatible_positive_pairs_per_rank metric during training. In each training process, TAO counts the promoted pairs that the loss of that process evaluates, and logs this count with synchronization across processes. The metric name uses the term rank for a training process. TAO does not log this metric for the ignore modes.

Requirements for Metadata-Aware Masking

Masking applies only with train.loss_type: siglip.

A violated condition raises a ValueError, except that an unsupported train.siglip_loss_dist_impl value raises a NotImplementedError. For every mode other than none, TAO checks the following conditions:

• You must set dataset.train.include_attribute_metadata to true.

• You must set dataset.train.type to custom.

• You must set train.siglip_loss_dist_impl to local or gather.

• Every entry under dataset.train.datasets must set image_list_file and train_pairs_file. The train_pairs_file must name an existing JSON file that has one entry for each non-blank line in image_list_file.

• Each entry of the pairs file must contain the keys image_attr_values and text_attr_values. If an entry contains an accessory key, it must contain both image_accessory_ids and text_accessory_ids.

• If you list more than one entry under dataset.train.datasets, or you use the pairs file for validation, an attribute_vocab.json file must sit next to each pairs file. With a single training entry, TAO logs a warning when the file is missing.

• With train.siglip_loss_dist_impl set to gather and more than one process, every process must have the same local batch size. Otherwise, TAO raises a RuntimeError.

For the positive modes only, TAO also checks the following conditions:

• You must set train.triplet_loss_weight to 0.0.

• You must list exactly one entry under dataset.train.datasets.

Refer to Dataset Configuration for the dataset fields.

optim

train:
  optim:
    optimizer_type: adamw
    vision_lr: 1.0e-4
    text_lr: 1.0e-4
    weight_decay: 1.0e-4
    betas: [0.9, 0.95]
    eps: 1.0e-6
    warmup_steps: 100
    scheduler: cosine

Field	Data type	Description	Default value	Valid options
optim.optimizer_type	string	Optimizer. Use lamb for large-batch distributed training.	adamw	adamw, lamb
optim.vision_lr	float	Learning rate for the vision encoder.	1.0e-4
optim.text_lr	float	Learning rate for the text encoder.	1.0e-4
optim.weight_decay	float	L2 regularization coefficient.	1.0e-4
optim.betas	list[float]	Adam/LAMB beta parameters.	[0.9, 0.95]
optim.eps	float	Epsilon for numerical stability.	1.0e-6
optim.warmup_steps	int	Number of linear warmup steps at the start of training.	100
optim.scheduler	string	Learning rate schedule after warmup.	cosine	cosine, constant, linear

peft

The peft group enables Parameter-Efficient Fine-Tuning (PEFT) with Low-Rank Adaptation (LoRA). In lora mode, TAO trains injected low-rank adapters on the linear layers whose leaf name is in target_modules, instead of the full tower.

peft:
  enabled: true
  method: lora
  train_logit_calibration: true
  vision:
    mode: lora
    target_modules: [q_proj, k_proj, v_proj, out_proj]
    num_last_blocks: 3
    rank: 8
    alpha: 16
    dropout: 0.05
  text:
    mode: frozen
    target_modules: [q_proj, k_proj, v_proj, out_proj]
    num_last_blocks: 3
    rank: 8
    alpha: 16
    dropout: 0.05

The default value of peft.enabled is false, and the default mode of each tower is frozen. The example turns on PEFT and adapts the vision tower with LoRA. All other values in the example are the default values.

For TAO 7.2 specifications, replace tower-level enabled: true with mode: lora and enabled: false with mode: frozen; remove the old tower-level keys. Keep peft.enabled. The new peft.train_logit_calibration defaults to true.

When peft.enabled is true, TAO first freezes every model parameter. TAO then applies the mode of each tower:

• frozen: TAO keeps the tower frozen.

• full: TAO makes every parameter of the tower trainable.

• lora: TAO wraps the matching linear layers in the last num_last_blocks blocks of the tower with LoRA adapters. Only the injected adapter parameters of that tower are trainable.

When you enable PEFT, the tower mode overrides model.freeze_vision_encoder and model.freeze_text_encoder. TAO builds a tower as frozen only when its mode is frozen. TAO logs a warning for a tower in the following cases:

• The freeze flag of the tower is true and its mode is not frozen.

• The freeze flag of the tower is false and its mode is frozen.

With the default values (false freeze flags and frozen modes), TAO logs the warning for each tower. Refer to Model Configuration for the freeze flags.

The train_logit_calibration field controls the logit parameters independently of the tower modes. When it is true and you enable PEFT, TAO trains logit_scale and, when the model has one, logit_bias. With both towers in frozen mode and train_logit_calibration: true, only the logit parameters train. If no parameter is trainable after TAO applies the tower modes, for example when both towers are frozen and train_logit_calibration is false, TAO raises a ValueError.

The method field has the single valid option lora.

These rules apply to target_modules and num_last_blocks:

• TAO matches each name in target_modules against the exact leaf module name (the text after the last dot) of the linear layers inside the last num_last_blocks blocks of the tower. TAO skips a name that matches a module that is not a linear layer.

• Setting num_last_blocks: 0 adapts all blocks of the tower. For a tower in lora mode, a value outside the range from 0 to the number of blocks in the tower raises a ValueError.

• If no name in target_modules matches a linear layer in the targeted blocks, TAO raises a ValueError that lists the available leaf names of the linear layers in those blocks.

The default target_modules are the SigLIP2 attention projection names q_proj, k_proj, v_proj, and out_proj.

Note

LoRA mode does not support towers that use nn.MultiheadAttention, such as OpenCLIP. If any nn.MultiheadAttention module exists in the targeted blocks of a tower in lora mode, TAO raises a NotImplementedError. The full and frozen modes work for these towers.

TAO scales the output of each adapter by alpha divided by rank. TAO stores this scale value in the checkpoint together with the adapter weights. When you load a checkpoint that has no stored scale values, TAO derives the scale from the current rank and alpha values and logs a warning. The warning states that the rank and alpha values must match the values you used for training.

TAO validates the following values only for towers in lora mode, and raises a ValueError when a value is invalid:

• rank must be an integer that is at least 1.

• alpha must be positive and finite.

• dropout must be at least 0.0 and less than 1.0. The maximum in the field definition is 1.0, but TAO rejects dropout values of 1.0 or higher.

Field	Data type	Description	Default value	Valid options
enabled	bool	Whether to enable PEFT mode. When false, training uses standard fine-tuning and respects model.freeze_vision_encoder and model.freeze_text_encoder.	false
method	categorical	PEFT method. The only supported method is lora.	lora	lora
train_logit_calibration	bool	Whether to train logit_scale and, when the model has one, logit_bias while PEFT is enabled. When false, TAO keeps these parameters frozen.	true
vision	collection	LoRA configuration for the vision encoder.	Vision and Text Towers
text	collection	LoRA configuration for the text encoder.	Vision and Text Towers

Vision and Text Towers

peft.vision and peft.text take the same set of fields, so one table covers both towers. The default values are the same for both towers.

Field	Data type	Description	Default value	Valid options
mode	categorical	Tower adaptation mode. frozen leaves the tower fixed; full trains all tower parameters; lora trains only injected LoRA parameters.	frozen	frozen, full, lora
target_modules	list	Module leaf names to target for LoRA injection. SigLIP2: q_proj, k_proj, v_proj, out_proj. Radio-CLIP: qkv, proj (fused attention). OpenCLIP uses nn.MultiheadAttention and does not support lora mode; use full or frozen.	['q_proj', 'k_proj', 'v_proj', 'out_proj']
num_last_blocks	int	Number of final transformer blocks to adapt. 0 means adapt all blocks.	3	min 0
rank	int	LoRA rank (low-rank dimension).	8	min 1
alpha	int	LoRA alpha scaling factor. The effective scale is alpha / rank.	16	min 1
dropout	float	Dropout applied to the LoRA input. For a tower in lora mode, the value must be less than 1.0; TAO raises an error for 1.0 or higher.	0.05	0.0 <= dropout < 1.0

regularization

The regularization group adds losses that constrain how far the embeddings of the model you train drift from the pretrained model during fine-tuning. The model you train is the student. The teacher is a frozen copy of the pretrained model.

regularization:
  enabled: true
  embedding_mse_weight: 0.05
  cosine_weight: 0.05
  similarity_weight: 0.10

The default value of regularization.enabled is false. The example turns regularization on and uses the default weights.

When regularization.enabled is true, TAO builds the teacher as a frozen deep copy of the model before it injects LoRA adapters, so the teacher keeps the original pretrained weights. TAO builds the teacher whenever you enable regularization, whether or not peft.enabled is true.

TAO adds the weighted sum of three terms to the contrastive loss:

loss = contrastive_loss
       + embedding_mse_weight * embedding_mse
       + cosine_weight * cosine
       + similarity_weight * similarity

The three terms are:

• embedding_mse is the mean of the mean squared error between the student and teacher image embeddings and the mean squared error between the student and teacher text embeddings.

• cosine is the mean of the image term and the text term, where each term is 1 minus the mean cosine similarity between the student and teacher embeddings.

• similarity is the mean squared error between the student and teacher image-text similarity matrices.

TAO excludes the teacher from checkpoints. TAO rebuilds the teacher from the base weights each time it constructs the model.

When you enable regularization, TAO logs the unweighted value of each term as a training metric:

• train/contrastive_loss

• train/emb_mse_loss

• train/cosine_loss

• train/sim_loss

When you configure a validation retrieval dataset, TAO also logs val/img_drift, val/txt_drift, and val/embedding_drift. Each drift value is 1 minus the mean cosine similarity between the student and teacher embeddings over the validation set, and val/embedding_drift is the mean of the image and text drift values.

Field	Data type	Description	Default value	Valid options
enabled	bool	Whether to enable preservation regularization. When false, TAO uses only the contrastive loss.	false
embedding_mse_weight	float	Weight for the mean squared error (MSE) loss between student and teacher embeddings.	0.05	min 0.0
cosine_weight	float	Weight for the cosine preservation loss between student and teacher embeddings.	0.05	min 0.0
similarity_weight	float	Weight for the similarity matrix preservation loss (MSE between student and teacher image-text similarity matrices).	0.1	min 0.0

dataset

dataset:
  seed: 42
  train:
    type: custom
    datasets:
      - image_dir: /data/train/images
        caption_dir: /data/train/captions
        caption_file_suffix: .txt
        image_list_file: null
        train_pairs_file: null
        attribute_pairs_file: null
    batch_size: 16
    num_workers: 8
    balance_query_types: false
    unique_caption_per_batch: true
    include_attribute_metadata: false
  val:
    datasets: []
    batch_size: 16
    num_workers: 8

Field	Data type	Description	Default value	Valid options
dataset.seed	int	Random seed for data loading and shuffling.	42
dataset.train.type	string	Dataset format for training.	custom	custom, wds
dataset.train.datasets	list	List of dataset entries. Each entry specifies image_dir, caption_dir, caption_file_suffix, image_list_file, train_pairs_file, and attribute_pairs_file. TAO concatenates multiple entries.	[]
dataset.train.datasets[].image_list_file	string	Path to a text file that lists image paths relative to image_dir, one per line. TAO skips blank lines and does not check that the listed files exist. If you do not set this field, TAO scans image_dir for .jpg, .jpeg, and .png files that have a caption file.	null
dataset.train.datasets[].train_pairs_file	string	Optional path to the pairs file of this dataset entry. TAO reads it for training when dataset.train.include_attribute_metadata or dataset.train.balance_query_types is true.	null
dataset.train.datasets[].attribute_pairs_file	string	Optional path to the pairs file that TAO uses for metadata-aware validation and for direct evaluation with evaluate.datasets. TAO does not read it for training. If you leave it unset and image_list_file ends in _list.txt, TAO replaces that suffix with _pairs.json for validation and direct evaluation. Entries of dataset.val.datasets use this field in the same way.	null
dataset.train.batch_size	int	Batch size per GPU during training.	16
dataset.train.num_workers	int	Number of dataloader worker processes.	8
dataset.train.balance_query_types	bool	Whether to balance CLIP training batches using the train_pairs_file metadata. For every training dataset entry, this requires image_list_file and an existing train_pairs_file. Otherwise, TAO disables balanced sampling and logs a warning. This field applies to the custom dataset type.	false
dataset.train.unique_caption_per_batch	bool	Whether to build batches in which no two rows share a caption string. This field takes effect only when balance_query_types is true. If TAO cannot set up unique-caption batching, it falls back to balancing by query type and logs a warning.	true
dataset.train.include_attribute_metadata	bool	Whether to include image and text attribute tensors in custom training batches. Required when train.loss_type is siglip and train.siglip_loss_mask_mode is not none. Every custom training dataset entry then needs image_list_file and train_pairs_file. Refer to Metadata-Aware SigLIP Loss Masking.	false

val

dataset:
  val:
    datasets:
      - image_dir: /data/val/images
        caption_dir: /data/val/captions
        caption_file_suffix: .txt
        image_list_file: /data/val/val_list.txt
    batch_size: 16
    num_workers: 8
    metadata_match_eval: false
    metadata_match_mode: scalar_attributes

The dataset.val group configures the validation datasets. Each entry of dataset.val.datasets takes the same fields as an entry of dataset.train.datasets. Refer to Dataset Configuration. When dataset.val.metadata_match_eval is true, TAO builds the validation dataloader with attribute metadata.

Field	Data type	Description	Default value	Valid options
dataset.val.datasets	list	List of validation dataset entries.	[]
dataset.val.batch_size	int	Batch size per GPU.	16	min 1
dataset.val.num_workers	int	Number of data-loading worker processes.	8	min 0
dataset.val.metadata_match_eval	bool	Whether to use attribute metadata to define the text-to-image validation positives. When true, validation logs metrics whose names start with val/pas/ instead of the paired retrieval metrics. Each validation dataset then needs an existing pairs file. TAO takes the file from attribute_pairs_file or infers it from image_list_file, as the dataset.train.datasets[].attribute_pairs_file row of the Dataset Configuration table describes. Multiple datasets must use identical attribute and accessory vocabularies. When false, validation keeps the paired diagonal ground truth.	false
dataset.val.metadata_match_mode	categorical	Metadata compatibility that TAO uses for the text-to-image validation positives when metadata_match_eval is true. With scalar_attributes, an image is a positive for a query when every attribute that the query specifies matches the image. With scalar_plus_accessories, the accessory IDs of the query must also appear in the accessory IDs of the image.	scalar_attributes	scalar_attributes, scalar_plus_accessories

augmentation

dataset:
  augmentation:
    scale: [0.4, 1.0]
    color_jitter: [0.8, 0.32, 0.32, 0.32, 0.08]
    grayscale: 0.2

Field	Data type	Description	Default value
augmentation.scale	list[float]	Random resized crop scale range [min, max]. Set to [1.0, 1.0] to disable cropping.	[0.4, 1.0]
augmentation.color_jitter	list[float]	Color jitter parameters: [probability, brightness, contrast, saturation, hue]. Set to [] to disable.	[0.8, 0.32, 0.32, 0.32, 0.08]
augmentation.grayscale	float	Probability of converting an image to grayscale during training.	0.2

evaluate

evaluate:
  checkpoint: /results/clip_experiment/train/epoch_100.pth
  batch_size: 16
  datasets: []
  pas_ground_truth_mode: paired_caption

Field	Data type	Description	Default value	Valid options
evaluate.checkpoint	string	Path to the TAO training checkpoint. When set to null, TAO loads pretrained weights directly — enabling zero-shot evaluation.	null
evaluate.batch_size	int	Batch size for embedding extraction during evaluation.	16
evaluate.datasets	list	List of dataset path configurations. When this list is non-empty, TAO runs a direct evaluation instead of retrieval over dataset.val.datasets. The direct evaluation embeds the images and queries of a pairs file, evaluates them, and writes the results as CSV files. TAO uses the pairs file of the first entry that has an existing pairs file. TAO takes that file from attribute_pairs_file or infers it from image_list_file, as the dataset.train.datasets[].attribute_pairs_file row of the Dataset Configuration table describes. TAO raises a ValueError when no entry has an existing pairs file. When the list is empty, TAO runs retrieval over dataset.val.datasets.	[]
evaluate.pas_ground_truth_mode	categorical	Ground-truth policy for direct evaluation. TAO reads this field only when evaluate.datasets is non-empty. The value paired_caption evaluates with exact-caption pairs only. With scalar_attributes or scalar_plus_accessories, TAO also evaluates with positives derived from the attributes. With scalar_plus_accessories, the accessory IDs of the query must also appear in the accessory IDs of the image.	paired_caption	paired_caption, scalar_attributes, scalar_plus_accessories

inference

inference:
  checkpoint: /results/clip_experiment/train/epoch_100.pth
  datasets:
    - image_dir: /data/inference/images
  text_file: /data/inference/prompts.txt
  batch_size: 16

Field	Data type	Description	Default value
inference.checkpoint	string	Path to the TAO training checkpoint. When set to null, TAO loads pretrained weights directly.	null
inference.datasets	list	List of image dataset entries. Each entry sets image_dir and, optionally, image_list_file. Without an image list, TAO scans image_dir recursively for .jpg, .jpeg, .png, .bmp, .gif, and .webp files. Refer to Using an Image List.	[]
inference.text_file	string	Path to a plain text file with one prompt per line. TAO extracts a text embedding for each prompt.	null
inference.batch_size	int	Batch size for embedding extraction.	16

export

export:
  checkpoint: /results/clip_experiment/train/epoch_100.pth
  onnx_file: /results/clip_experiment/export/clip_model.onnx
  encoder_type: combined
  input_height: 256
  input_width: 256
  batch_size: -1
  opset_version: 17

Field	Data type	Description	Default value
export.checkpoint	string	Path to the TAO training checkpoint. When set to null, TAO exports the pretrained model directly.	null
export.onnx_file	string	Output path for the ONNX file.
export.encoder_type	string	Controls whether TAO exports a single combined encoder or two separate encoders. Refer to Exporting the Model for guidance on which to choose.	combined
export.input_height	int	Input image height in pixels.	256
export.input_width	int	Input image width in pixels.	256
export.batch_size	int	Export batch size. Set to -1 for a dynamic batch axis.	-1
export.opset_version	int	ONNX opset version. The minimum supported value is 11.	17

Training the Model

Note

To run multi-GPU training, set train.num_gpus in the specification file or pass it as a command-line override. For multi-node training, set train.num_nodes and train.distributed_strategy: ddp. Use distributed_strategy: fsdp for models that exceed single-node GPU memory.

Evaluating the Model

CLIP evaluation runs bidirectional retrieval across your validation dataset and reports the following metrics:

• R@1, R@5, R@10: Recall at k. The fraction of queries for which the correct match appears in the top-k retrieved results.

• mAP: Mean average precision across all queries.

• Median Rank: The median rank position of the first correct match across all queries. Lower is better.

• Mean Rank: The mean rank position of the first correct match. Lower is better.

• AUC: Area under the precision-recall curve.

TAO reports all metrics for two directions: image-to-text (given an image, retrieve the matching caption) and text-to-image (given a caption, retrieve the matching image).

When evaluate.datasets is non-empty, TAO runs a direct evaluation instead of retrieval over dataset.val.datasets and writes the results as CSV files. Refer to Evaluate Configuration.

With dataset.val.metadata_match_eval: true, training logs val/pas/{easy,medium,hard,overall}_{mAP,rank1,rank5}. For comparison, use the same checkpoint and data, and match evaluate.pas_ground_truth_mode to dataset.val.metadata_match_mode. Read nvidia_pas_metadata_metrics_weighted_aggregate.csv; nvidia_pas_metrics_weighted_aggregate.csv retains paired-caption results. Direct PAS evaluation writes CSVs and does not emit test/t2i_mAP.

Running Inference

The inference task extracts image and text embeddings and saves them as HDF5 files in results_dir:

• image_embeddings.h5: Contains datasets embeddings (float32, shape N × D) and image_paths (string).

• text_embeddings.h5: Contains datasets embeddings (float32, shape N × D) and texts (string).

All embeddings are L2-normalized before saving.

For examples of how to use these embeddings in downstream applications, refer to Using CLIP Embeddings.

Using an Image List

By default, inference embeds every image that it finds under image_dir. TAO scans image_dir recursively for files with the extensions .jpg, .jpeg, .png, .bmp, .gif, and .webp, and processes them in sorted order. TAO matches the extensions without regard to case. When inference.datasets has several entries, TAO concatenates the results of the entries in the order that you list them.

To embed only selected images, set image_list_file on an entry of inference.datasets to a text file that lists the images.

inference:
  checkpoint: /results/clip_experiment/train/epoch_100.pth
  datasets:
    - image_dir: /data/inference/images
      image_list_file: /data/inference/image_list.txt
  batch_size: 16

The list file has one relative path per line, such as:

cats/img_001.jpg
dogs/img_014.png

When you set image_list_file, TAO applies the following rules:

• TAO embeds only the listed files, in list order.

• Each line is a path relative to image_dir. TAO skips blank lines.

• TAO does not filter the listed files by extension.

• TAO skips a listed image that it cannot load and logs a warning, so the output can contain fewer embeddings than the list has entries.

TAO raises a ValueError in the following cases:

• A line is an absolute path or contains a .. path component. The message is image_list_file must contain paths relative to image_dir: <path>. TAO validates every line before it reads any image.

• The operating system reports an error when TAO opens or reads the list file. The message is Cannot read image_list_file: <path>.

An empty list file selects no images and does not fall back to scanning the directory. If no entry of inference.datasets yields any image, TAO logs the warning No images found in any dataset and writes no image_embeddings.h5.

Exporting the Model

TAO exports CLIP models to ONNX. You can export either a single combined encoder or two separate encoders, depending on your deployment requirements.

Combined encoder (encoder_type: combined): Produces a single ONNX file containing both the vision and text encoders.

Direction	Details
Inputs	image (B × 3 × H × W, float32), input_ids (B × seq_len, int64), attention_mask (B × seq_len, int64)
Outputs	image_embedding (B × D), text_embedding (B × D), logit_scale (scalar), logit_bias (scalar)

Use this format when you run vision and text encoding together at inference time—for example, in real-time retrieval or classification pipelines.

Separate encoders (encoder_type: separate): Produces two ONNX files: clip_model_vision.onnx and clip_model_text.onnx.

Engine	Details
Vision	Input: image (B × 3 × H × W). Outputs: image_embedding, logit_scale, logit_bias.
Text	Inputs: input_ids, attention_mask. Output: text_embedding, logit_scale, logit_bias.

Use this format when you want to pre-compute text embeddings offline—for example, to index a fixed set of class names or captions once and then run only the vision encoder at query time.

Note

attention_mask is a required ONNX graph input but its values are not used—the model always substitutes an all-ones mask internally. Passing the tokenizer’s mask or np.ones_like(input_ids) produces identical results. Refer to Usage Notes for ONNX and TensorRT Deployment for details.

Warning

Currently, attention_mask is accepted as an explicit graph input for backward compatibility only. This input is deprecated and scheduled for removal. Remove it from your inference pipeline to avoid a future breaking change.

Note

clip export merges LoRA adapters into the base linear weights before it writes the ONNX file, so the exported model contains no LoRA modules. The merge runs when the model that TAO builds from the export specification contains LoRA modules. This requires peft.enabled: true with a tower in lora mode in the export specification. Refer to PEFT Configuration. When the model contains no LoRA modules, TAO merges nothing.

The export command also produces two artifact files alongside the ONNX output:

• <name>_config.yaml: Saved experiment configuration, required for engine generation and inference.

• <name>_tokenizer/: Saved HuggingFace tokenizer directory, required for text preprocessing.

Important

Keep _config.yaml and _tokenizer/ in the same directory as the ONNX file. Moving or renaming them causes engine generation and inference to fail.

Important

For models larger than 2 GB, ONNX export writes two files: the .onnx file and an external data file that stores the large weight tensors (an ONNX external data limitation). The external data file name is set by the ONNX export path configuration. Do not rename it after export; the .onnx file references it by the exact name written at export time. Both files must remain in the same directory. If you move the .onnx file, move the external data file alongside it, or the engine build cannot succeed.

Warning

siglip2-so400m-patch16-naflex cannot be exported to ONNX. Use a fixed-resolution variant such as siglip2-so400m-patch16-384 instead.

Usage Notes for ONNX and TensorRT Deployment

The following notes apply when you load the exported ONNX model or TRT engine directly.

Attention Mask Behavior

attention_mask is present as an ONNX graph input for backward compatibility, but its values are not used. The model always substitutes an all-ones mask internally. You can safely pass the tokenizer’s attention_mask or np.ones_like(input_ids); both produce identical results.

Warning

Currently, attention_mask is accepted as an explicit graph input for backward compatibility only. This input is deprecated and scheduled for removal. Remove it from your inference pipeline to avoid a future breaking change.

Sequence Length

The text inputs (input_ids and attention_mask) must use the same max_length passed to the tokenizer. CLIP tokenizers typically use 77; SigLIP2 tokenizers use 64. Passing a different length causes a shape mismatch at runtime.

Dynamic Batch and TensorRT Shape Profiles

When export.batch_size: -1, the batch dimension is dynamic. When building a TRT engine with trtexec, provide --minShapes, --optShapes, and --maxShapes for every input. For a combined encoder with 77-token sequences:

trtexec --onnx=clip_model.onnx \
  --minShapes=image:1x3x256x256,input_ids:1x77,attention_mask:1x77 \
  --optShapes=image:8x3x256x256,input_ids:8x77,attention_mask:8x77 \
  --maxShapes=image:32x3x256x256,input_ids:32x77,attention_mask:32x77

The attention_mask shape profile must match input_ids exactly. After the deprecation takes effect, omit attention_mask from all three shape arguments.

Image Preprocessing

The image tensor must be preprocessed to the same pixel statistics used during training. When loading the ONNX model directly, apply the per-channel mean and standard deviation stored in <name>_config.yaml (exported alongside the ONNX file).

Logit Scale and Logit Bias

Both logit_scale and logit_bias are exported as scalar outputs. For SigLIP-style models, compute the match probability as sigmoid(logit_scale * dot(image_emb, text_emb) + logit_bias). For CLIP-style models, logit_bias is zero and softmax over class scores works as well. When using separate encoders, logit_scale and logit_bias are available from either encoder; you only need one copy.