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.	false	true, false
model.freeze_text_encoder	bool	Freeze text encoder weights during training.	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
  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.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

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

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
    batch_size: 16
    num_workers: 8
  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, and image_list_file. Multiple entries are concatenated.
dataset.train.batch_size	int	Batch size per GPU during training.	16
dataset.train.num_workers	int	Number of dataloader worker processes.	8

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

Field	Data type	Description	Default value
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

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. Supported image extensions: .jpg, .jpeg, .png, .bmp, .gif, .webp.
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).

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.

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.

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.