tao-train-nvdinov2

NVDINOv2

NVDINOv2 用于自监督视觉表征学习。通过自蒸馏（教师-学生）训练视觉 Transformer，无需标签。生成通用视觉特征。

为预训练 ViT 权重设置 train.pretrained_model_path。

对于 TAO Deploy TensorRT 操作（gen_trt_engine），请先阅读 references/tao-deploy-nvdinov2.md。部署规范模板位于此技能的 references/ 文件夹中，文件名前缀为 spec_template_deploy_*.yaml。

Dataclass 架构

生成的 TAO Core 架构打包在 schemas/<action>.schema.json 中，schemas/manifest.json 列出了可用操作。每个生成的架构还会根据架构顶层的 default 字段生成 references/spec_template_<action>.yaml。AutoML 启用状态通过 references/skill_info.yaml 中模型层级的 automl_enabled 声明。要运行 AutoML，仍需存在且可解析的 schemas/train.schema.json 和 references/spec_template_train.yaml。请使用打包的训练架构获取 automl_default_parameters、automl_disabled_parameters、默认值、最小值/最大值边界、枚举、选项权重、数学条件、依赖项和常用参数。不要期望运行时存在 ~/tao-core；维护者会在打包技能库前重新生成架构/模板。

训练操作策略

此模型在模型层级启用了 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。保留工作流/应用对数据集、规范、输出目录、GPU/平台设置、父检查点和 automl_policy 的覆盖。仅当 automl_policy: off 或缺少已打包的训练架构/模板时，才直接进行模型训练；若缺少架构，请说明 AutoML 已启用，但在为此模型生成架构前无法运行。

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

训练要求

• 数据集类型： image_classification
• 格式： ssl
• 监控指标： train_loss

各操作的数据集要求

操作	规范键	来源	文件	是否为列表？
distill	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	否
inference	dataset.test_dataset.images_dir	inference_dataset	images_test.tar.gz	否
train	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	否

典型规范覆盖项

每个操作都必须覆盖数据源——Agent 必须根据上表中的各操作数据集要求构造数据源路径，并将其纳入 spec_overrides。

S3_TRAIN = "s3://bucket/data/train"
S3_EVAL = "s3://bucket/data/eval"

train（必需数据源）：

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
}

distill（必需数据源）：

{
    "dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
}

inference（必需数据源）：

{
    "dataset.test_dataset.images_dir": f"{S3_EVAL}/images_test.tar.gz",
}

评估数据集

可选。SSL 训练不使用标签。评估取决于下游任务。

重要参数

• model.backbone.teacher_type：教师 ViT 变体。默认值为 vit_l（ViT-Large）。
• model.backbone.student_type：学生 ViT 变体。默认值为 vit_l。通常与教师模型一致。
• model.backbone.img_size：输入图像大小。默认值为 518。更高的分辨率可生成更好的特征，但会占用更多内存。
• model.backbone.patch_size：ViT 图像块大小。默认值为 14。
• dataset.batch_size：每个 GPU 的批次大小。默认值为 4。由于需要分别执行教师和学生的前向传播，SSL 训练的内存占用较高。
• train.layerwise_decay：逐层学习率衰减。对 ViT 微调很重要。
• train.clip_grad_norm：梯度裁剪。对稳定 SSL 训练很重要。

多 GPU / 多节点

**启动方式：**由 Lightning 管理（单个 python 进程，由 Lightning 启动工作进程）。

Spec 键	说明	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]
train.num_nodes	节点数量	1

• 策略：auto（Lightning 自动选择最佳策略）
• 始终启用 sync_batchnorm——这对使用教师-学生框架进行 SSL 训练至关重要
• 强烈建议使用多 GPU（4-8 个 GPU）进行有意义的 SSL 训练

多节点环境变量（由编排器设置）：WORLD_SIZE、NODE_RANK、MASTER_ADDR、MASTER_PORT、NUM_GPU_PER_NODE。

硬件

至少需要 4 个 GPU，建议使用 8 个 GPU。每个 GPU 需配备 40GB 以上显存（建议使用 A100）。使用 ViT-Large 教师和学生模型进行 SSL 的内存占用非常高。需要使用显存为 40GB 以上的 A100 GPU。强烈建议使用多 GPU。

错误模式

CUDA 内存不足：使用 img_size=518 的 ViT-Large 教师和学生模型需要 40GB 以上的 GPU 显存。请减小 batch_size、img_size，或使用较小的 ViT 变体。

收敛缓慢：SSL 需要许多轮次。默认值 10 适用于快速测试；生产运行通常使用 100 轮以上。

Spec 参数 / 父模型推断

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

来自 TAO Core nvdinov2.config.json 的推断映射：

操作	Spec 字段	推断函数	含义
distill	encryption_key	key	加密密钥
distill	model.distill.pretrained_non_distill_pl_model_path	parent_model	从父作业结果文件夹推断出的模型文件
distill	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	当前作业的结果目录
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，也不要修改生成的运行器脚本来猜测检查点路径。

部署

• tao-deploy-nvdinov2 — 使用 TAO Deploy 生成 TensorRT 引擎的 NvDINOv2 部署工作流。