tao-finetune-clip

CLIP

用于零样本和微调图像分类、图文检索及嵌入提取的对比式语言-图像预训练模型。微调可使 CLIP 的共享图像-文本嵌入空间适配特定领域的图像-描述数据。

无需默认的 NGC 预训练检查点。如果未设置 train.pretrained_model_path、evaluate.checkpoint、inference.checkpoint 或 export.checkpoint，TAO 会从 HuggingFace 加载 SigLIP2/OpenCLIP 变体的预训练权重，或通过 torch.hub 加载 Radio-CLIP 的预训练权重，因此首次使用时需要网络访问权限或本地镜像。

支持的操作：train、evaluate、inference、export、gen_trt_engine。

训练操作策略

此模型在模型层启用了 AutoML。处理任何训练阶段请求之前，请先读取 references/skill_info.yaml，并根据明确指定的 automl_policy 值或用户的工作流请求确定本次运行的覆盖设置。将“关闭 AutoML”“禁用 AutoML”“不使用 HPO”或“普通训练”等说法视为仅对本次运行设置 automl_policy: off；否则默认为 auto。当 automl_policy: auto、automl_enabled: true，且已打包 schemas/train.schema.json 和 references/spec_template_train.yaml 时，默认通过 tao-skill-bank:tao-run-automl 执行训练操作，并传入此模型的 skill_dir。保留工作流/应用的覆盖设置，包括数据集、spec、输出目录、GPU/平台设置、父检查点和 automl_policy。仅当 automl_policy: off 或未打包训练 schema/template 时，才直接训练模型；若缺少 schema，请说明 AutoML 已启用，但在生成 schema 之前无法用于此模型。

evaluate、inference、export 和部署流程等非训练操作仍由此模型技能处理。每次运行的 automl_policy 覆盖设置不会更改模型元数据。

说明

使用此技能执行 NVIDIA TAO CLIP 任务：训练、评估、嵌入推理、ONNX 导出和 TensorRT 引擎生成。首先确定请求的操作，然后仅加载该操作所需的引用文件：默认参数使用 defaults.json，操作/数据源配置使用 config.json，完整 spec 结构使用 references/spec_template.yaml，SDK 元数据使用 references/model_info.yaml。

对于需要数据集的操作，请收集用户提供的图像、描述、列表或提示词文件，并将解析后的路径放入 spec_overrides。对于 export 和 gen_trt_engine，如果上游任务可用，则推断父级产物；否则要求用户明确提供检查点、ONNX 或引擎路径。在 TAO Deploy 镜像中运行 gen_trt_engine、TensorRT evaluate 和 TensorRT inference。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请先阅读 references/tao-deploy-clip.md。部署 spec 模板位于此技能的 references/ 文件夹中，文件名前缀为 spec_template_deploy_。

训练要求

• 数据集类型： image_text
• 格式： 自定义图像/描述文件夹或 WebDataset 分片
• 监控指标： val/t2i_mAP

支持的模型

• SigLIP2： siglip2-so400m-patch16-256（默认）、siglip2-so400m-patch14-224、siglip2-so400m-patch14-384、siglip2-so400m-patch16-384、siglip2-so400m-patch16-512、siglip2-so400m-patch16-naflex
• Radio-CLIP： c-radio_v3-b、c-radio_v3-l、c-radio_v3-h、c-radio_v3-g
• OpenCLIP / NV-CLIP： ViT-L-14-SigLIP-CLIPA-224、ViT-L-14-SigLIP-CLIPA-336、ViT-H-14-SigLIP-CLIPA-224、ViT-H-14-SigLIP-CLIPA-336、ViT-H-14-SigLIP-CLIPA-574

Radio-CLIP 需要将 model.adaptor_name 设置为 siglip 或 clip。

各操作的数据集要求

操作	规格键	来源	文件	列表？
train	dataset.train.datasets	train_datasets	image_dir: images.tar.gz, image_list_file: image_list.txt, caption_dir: captions.tar.gz	是
train	dataset.train.wds.root_dir	train_wds_dataset	包含 .tar 分片的根目录	否
train	dataset.train.wds.shard_list_file	train_wds_dataset	列出分片路径的 shards.txt	否
train	dataset.val.datasets	eval_dataset	image_dir: images.tar.gz, image_list_file: image_list.txt, caption_dir: captions.tar.gz	是
evaluate	dataset.val.datasets	eval_dataset	image_dir: images.tar.gz, image_list_file: image_list.txt, caption_dir: captions.tar.gz	是
inference	inference.datasets	inference_dataset	image_dir: images.tar.gz	是
inference	inference.text_file	inference_dataset	prompts.txt	否
export	export.checkpoint	父级 train 作业或显式指定的检查点	checkpoint .pth；使用预训练权重导出时可选	否
gen_trt_engine	gen_trt_engine.onnx_file	父级 export 作业或显式指定的 ONNX	clip_model.onnx	否

对于自定义训练，将 dataset.train.type: custom，并提供 dataset.train.datasets 条目。图像文件和字幕文件必须具有相同的基本名称。caption_file_suffix 默认为 .txt，image_list_file 为可选项。

对于 WDS 训练，将 dataset.train.type: wds，并至少提供 dataset.train.wds.root_dir 或 dataset.train.wds.shard_list_file 其中一项。系统会递归扫描 root_dir 以查找 .tar 分片。shard_list_file 是一个文本文件，每行列出一个分片路径；相对路径会基于列表文件所在目录解析，如果同时提供了 root_dir，则基于 root_dir 解析。验证/评估数据仍使用 dataset.val.datasets 指定的自定义格式。

常见规格覆盖项

对于基于数据集的操作，必须覆盖数据源。根据“各操作的数据集要求”表构造路径，并将其加入 spec_overrides。对于 inference，至少提供 inference.datasets 或 inference.text_file 其中一项。

S3_TRAIN = "s3://bucket/data/train"
S3_WDS = "s3://bucket/data/wds"
S3_EVAL = "s3://bucket/data/eval"
S3_INFER = "s3://bucket/data/infer"

train，自定义数据集：

{
    "train.num_epochs": 10,
    "dataset.train.type": "custom",
    "dataset.train.datasets": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "image_list_file": f"{S3_TRAIN}/image_list.txt", "caption_dir": f"{S3_TRAIN}/captions.tar.gz"}],
    "dataset.val.datasets": [{"image_dir": f"{S3_EVAL}/images.tar.gz", "image_list_file": f"{S3_EVAL}/image_list.txt", "caption_dir": f"{S3_EVAL}/captions.tar.gz"}],
}

train，WDS 数据集：

{
    "train.num_epochs": 10,
    "dataset.train.type": "wds",
    "dataset.train.wds.root_dir": f"{S3_WDS}",
    "dataset.train.wds.shard_list_file": f"{S3_WDS}/shards.txt",
    "dataset.train.wds.samples_per_shard": 10000,
    "dataset.val.datasets": [{"image_dir": f"{S3_EVAL}/images.tar.gz", "image_list_file": f"{S3_EVAL}/image_list.txt", "caption_dir": f"{S3_EVAL}/captions.tar.gz"}],
}

evaluate：

{
    "dataset.val.datasets": [{"image_dir": f"{S3_EVAL}/images.tar.gz", "image_list_file": f"{S3_EVAL}/image_list.txt", "caption_dir": f"{S3_EVAL}/captions.tar.gz"}],
}

使用预训练权重进行零样本评估时，不设置 evaluate.checkpoint。对于 TensorRT 评估，设置 evaluate.trt_engine，而不是 evaluate.checkpoint。

inference：

{
    "inference.datasets": [{"image_dir": f"{S3_INFER}/images.tar.gz"}],
    "inference.text_file": f"{S3_INFER}/prompts.txt",
}

Inference 会在 results_dir 下写入 image_embeddings.h5 和/或 text_embeddings.h5。保存的嵌入经过 L2 归一化。

export：

{
    "export.onnx_file": "${results_dir}/export/clip_model.onnx",
    "export.encoder_type": "combined",
    "export.batch_size": -1,
}

如果部署时需要使用独立的视觉编码器和文本编码器，请设置 export.encoder_type: separate。独立导出会根据 export.onnx_file 的基础路径生成 _vision.onnx 和 _text.onnx 变体。

gen_trt_engine：

{
    "gen_trt_engine.onnx_file": "${results_dir}/export/clip_model.onnx",
    "gen_trt_engine.trt_engine": "${results_dir}/deploy/clip_model.engine",
    "gen_trt_engine.batch_size": -1,
    "gen_trt_engine.tensorrt.data_type": "fp16",
    "gen_trt_engine.tensorrt.min_batch_size": 1,
    "gen_trt_engine.tensorrt.opt_batch_size": 1,
    "gen_trt_engine.tensorrt.max_batch_size": 16,
}

评估数据集

训练时可选。如果提供了数据集，系统会在验证间隔计算验证指标。执行 evaluate 时则必须提供。

部署工作流

该技能将 gen_trt_engine 作为部署操作公开。在生成的 SDK 运行器中，使用 model_info["actions"]["gen_trt_engine"]，并在 TAO Deploy 镜像中运行，而不是在 PyTorch 训练镜像中运行。容器内的命令为 clip gen_trt_engine -e {config_path}；直接使用 TAO Launcher 时，同一操作的命令写作 tao deploy clip gen_trt_engine -e /path/to/spec.yaml。

TAO Deploy 同时支持组合式和独立编码器格式。对于独立编码器，将不带 _vision 或 _text 的基础路径传入 gen_trt_engine.onnx_file 和 gen_trt_engine.trt_engine；TAO 会自动检测或写入带后缀的视觉/文本文件。

使用 evaluate.trt_engine 进行 TensorRT 评估，使用 inference.trt_engine 提取 TensorRT 嵌入。这些 TensorRT 路径也需要在 TAO Deploy 镜像中运行。直接使用 TAO Launcher 时，这些命令分别写作 tao deploy clip evaluate 和 tao deploy clip inference。

TAO Deploy 完整参考：tao-deploy-clip。

重要参数

• model.type：主干网络系列和分辨率。部署时请使用固定分辨率的 SigLIP2/OpenCLIP 变体。
• model.adaptor_name：Radio-CLIP 必填。设置为 siglip 或 clip。
• model.image_size：训练变换使用的图像分辨率。请与所选固定分辨率主干网络保持一致。
• train.num_epochs：CLIP 微调通常能快速收敛。领域适配可先从 10-20 个 epoch 开始，只有在验证损失仍持续改善时才增加。
• train.optim.vision_lr / train.optim.text_lr：两个编码器的学习率。CLIP 对较高学习率较敏感；如果损失不稳定，请同时降低这两个学习率。
• model.freeze_vision_encoder / model.freeze_text_encoder：默认值为 false。数据集较小或只需适配一种模态时，冻结其中一个编码器可能会有所帮助。
• train.loss_type：SigLIP2 和 Radio-CLIP 推荐使用 siglip。CLIP 风格的 softmax 损失请使用 clip。
• export.encoder_type：combined 会导出一个 ONNX 图。separate 会分别导出视觉和文本图。
• gen_trt_engine.tensorrt.data_type：TensorRT 部署支持 fp16 和 fp32。

硬件

小型数据集可使用单 GPU 训练。对于超过 100k 张图像的数据集或大型主干网络，请使用 4+ 个 GPU。小型/固定分辨率运行建议每个 GPU 配备 16GB+ 显存；Radio-CLIP 或高分辨率 OpenCLIP 变体则需要更大的 GPU。

错误模式

CUDA 显存不足：减小 dataset.train.batch_size、dataset.val.batch_size 或 TensorRT 的 opt/max 批次大小。对于导出/部署，请检查 export.input_height 和 export.input_width 是否与所选固定分辨率主干网络一致。

损失为 NaN：微调时学习率过高。请降低 train.optim.vision_lr 和 train.optim.text_lr，增大 train.optim.warmup_steps，并确认字幕是有效的非空文本。

检索或分类效果为零：检查字幕和提示词是否与目标标签词汇表匹配。CLIP 会比较图像和文本嵌入，因此提示词措辞很重要。

数据集大小小于总批次大小：总批次大小为 batch_size * num_gpus。如果数据集（尤其是验证集）的样本数少于该值，请减小 dataset.val.batch_size 或 dataset.train.batch_size。

Radio-CLIP 配置验证错误：将 model.adaptor_name 显式设置为 siglip 或 clip。

Naflex 导出失败：当前 TAO 文档中，siglip2-so400m-patch16-naflex 仅用于训练，无法导出为 ONNX 或 TensorRT。请使用固定分辨率变体，例如 siglip2-so400m-patch16-384。

ONNX 外部数据缺失：大于 2 GB 的模型会导出一个 ONNX 文件和一个外部数据文件。请将这两个文件保存在同一目录中，并且在运行 gen_trt_engine 前不要重命名外部数据文件。

TensorRT 形状不匹配：使用动态批次导出时，请为每个输入提供 min/opt/max 形状配置。文本序列长度必须与分词器长度一致；CLIP 分词器通常为 77，SigLIP2 分词器通常为 64。

attention_mask 警告：为保持兼容性，当前导出的图会接受 attention_mask，但 TAO 会忽略其值，并可能在未来版本中将其移除。请勿构建依赖掩码值的新直接 ONNX 推理代码。

将 spec.yaml 与 schema 合并时出错：这是 Hydra/OmegaConf 配置验证错误。常见原因包括将 num_epochs 或 num_gpus 放在 spec 根目录，而非 train.* 下；或混淆训练图像尺寸（model.image_size）与导出尺寸（export.input_height 和 export.input_width）。

Spec 参数 / 父模型推断

特定于模型的推断映射应放在此 MD 文件中，而不是 config.json 中。生成的运行器应读取本节，并在调用 create_job() 前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程一致。

TAO Core clip.config.json 中的推断映射：

操作	Spec 字段	推断函数	含义
evaluate	encryption_key	key	加密密钥
evaluate	evaluate.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
evaluate	evaluate.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
evaluate	results_dir	output_dir	当前作业的结果目录
export	encryption_key	key	加密密钥
export	export.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
export	export.onnx_file	create_onnx_file	输出 ONNX 路径
export	results_dir	output_dir	当前作业的结果目录
gen_trt_engine	encryption_key	key	加密密钥
gen_trt_engine	gen_trt_engine.onnx_file	parent_model	从父作业结果文件夹推断出的模型文件
gen_trt_engine	gen_trt_engine.trt_engine	create_engine_file	输出 TensorRT 引擎路径
gen_trt_engine	results_dir	output_dir	当前作业的结果目录
inference	encryption_key	key	加密密钥
inference	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
inference	inference.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
inference	results_dir	output_dir	当前作业的结果目录
train	encryption_key	key	加密密钥
train	results_dir	output_dir	当前作业的结果目录
train	train.pretrained_model_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
train	train.resume_training_checkpoint_path	resume_model	从当前作业结果文件夹推断出的模型文件

对于 parent_model 或 parent_model_folder，请将上游 train/export/AutoML 子作业 ID 作为 parent_job_id 传入。SDK 会列出父作业的结果文件夹，筛选检查点资源，并返回选定的模型文件或文件夹。不要将这些映射加回 config.json，也不要修改生成的运行器脚本来猜测检查点路径。