tao-train-nvdinov2

NVDINOv2

单独安装？ 如果此会话并非由 TAO skill bank 插件初始化，请先运行 tao-setup 技能（主机预检、凭据、跨技能发现）。

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

设置 train.pretrained_model_path 以指定预训练 ViT 权重。

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

数据类 Schema

生成的 TAO Core Schema 打包在 schemas/<action>.schema.json 中，schemas/manifest.json 列出了可用操作。每个生成的 Schema 还会根据 Schema 顶层的 default 字段生成 references/spec_template_<action>.yaml。模型层的 AutoML 启用状态通过 references/skill_info.yaml 中的 automl_enabled 声明。某个操作要运行 AutoML，必须存在并能解析 schemas/<action>.schema.json 和 references/spec_template_<action>.yaml。使用已打包的所选操作 Schema 获取 automl_default_parameters、automl_disabled_parameters、默认值、最小值/最大值边界、枚举值、选项权重、数学条件、依赖项和常用参数。运行时不要假设存在 ~/tao-core；维护者会在打包 skill bank 前重新生成 Schema 和模板。

训练操作策略

此模型在模型层启用了 AutoML。处理任何训练阶段请求前，请阅读 references/skill_info.yaml，并根据明确指定的 automl_policy 值或用户的工作流请求确定本次运行的覆盖设置。默认使用 automl_policy: on，且在新的启动提示中仅提供 on / off 选项。将“turn off AutoML”“disable AutoML”“no HPO”或“plain training”等表述视为仅对本次运行设置 automl_policy: off。当 automl_policy: on、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 或缺少已打包的训练 Schema/模板时，才直接训练模型；若缺少 Schema，请说明此模型已启用 AutoML，但在生成 Schema 前无法运行 AutoML。

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

训练要求

• 数据集类型： image_classification
• 格式： ssl
• 监控指标： train_loss
• AutoML 指标约定： 使用训练操作输出的 train_loss， 并将其最小化。此模型未打包 evaluate 操作，因此不要自行设定 评估 KPI；请通过推理验证选定的 student_epoch_* 检查点。

各操作的数据集要求

操作	规格键	来源	源制品 → 运行时值	列表？
inference	dataset.test_dataset.images_dir	inference_dataset	images_test.tar.gz → 解压后的图像文件夹	否
train	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz → 解压后的图像文件夹	否

托管运行器可能会将这些数据集作为归档文件提供，因为 references/skill_info.yaml 声明了 runtime: extracted_folder。它必须 解压归档文件，并将生成的目录放入 images_dir。对于直接运行 Docker 或 TAO CLI 的情况，请在启动前解压归档文件；切勿将 .tar、.tar.gz、.tgz 或 .zip 文件设置为 images_dir 字段的值。

典型的 Spec 覆盖项

数据源覆盖项对于训练和推理均为必需——Agent 必须根据上方的“各操作数据集要求”表构建数据源路径，并将其包含在 spec_overrides 中。

TRAIN_IMAGES_DIR = "/workspace/data/extracted/train/images_train"
TEST_IMAGES_DIR = "/workspace/data/extracted/test/images_test"

train（必需的数据源）：

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
}

本地 AutoML 验证 / 冒烟测试： 如果目标是在本地 Docker 上确认贝叶斯启动、指标选择、最佳模型选择和检查点持久化，请使用此格式。它能让运行具有代表性，同时避免速度慢得多的 ViT-Large 默认设置。

{
    "wandb.enable": False,
    "model.backbone.teacher_type": "vit_s",
    "model.backbone.student_type": "vit_s",
    "model.backbone.img_size": 224,
    "dataset.batch_size": 8,
    "dataset.workers": 2,
    "train.num_epochs": 1,
    "train.checkpoint_interval": 1,
    "train.num_prototypes": 1024,
    "train.precision": "32-true",
    "train.use_custom_attention": False,
    "train.num_gpus": 1,
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
}

export（必需的检查点交接）：

{
    "export.checkpoint": "<selected train/AutoML student_epoch_* checkpoint>",
    "export.onnx_file": "/path/to/results/nvdinov2.onnx",
}

inference（必需的数据源）：

{
    "inference.checkpoint": "<selected train/AutoML student_epoch_* checkpoint>",
    "model.backbone.teacher_type": "<same value used for train>",
    "model.backbone.student_type": "<same value used for train>",
    "model.backbone.img_size": "<same value used for train>",
    "train.use_custom_attention": "<same value used for train>",
    # NVDINOv2 enables persistent_workers in its prediction DataLoader.
    # Keep this positive; dataset.workers=0 makes inference fail before loading data.
    "dataset.workers": 2,
    "dataset.test_dataset.images_dir": TEST_IMAGES_DIR,
}

评估数据集

可选。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 Key	说明	默认值
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。

错误模式

未找到图像／归档路径被拒绝：dataset.*.images_dir 必须是包含图像的已解压目录，而不是归档文件。请先解压 .tar、.tar.gz、.tgz 或 .zip 输入文件。由归档文件提供支持的托管数据源必须使用 runtime: extracted_folder 生成的目录。

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

推理检查点包含意外的 Lightning 键：进行下游 inference 时，请传入所选 AutoML 运行的 student_epoch_*.pth 检查点，而不是 nvdinov2_model_latest.pth。latest 文件是训练检查点，推理加载器会报告 state_dict、优化器状态和调度器状态等意外键。

推理失败，提示 persistent_workers option needs num_workers > 0： 请将 dataset.workers 设为正数（最小化本地运行时可使用 2）。 NVDINOv2 预测 DataLoader 会启用持久化工作进程，因此无法使用通用的零工作进程冒烟测试覆盖设置运行。

导出检查点包含意外的 Lightning 键：导出同样使用所选的 student_epoch_*.pth 检查点。仅在通过 train.resume_training_checkpoint_path 恢复／重新训练时，才使用完整的 model_epoch_*.pth 检查点。

将 TensorRT 引擎传给 PyT 推理：打包的 PyT nvdinov2 inference 实现只会加载 .pth 或 .tlt 模型路径。TAO Deploy gen_trt_engine 会为下游使用方构建 TensorRT 引擎，但 PyT 推理操作无法在该引擎上运行。

无法使用单独的 distill 操作：当前 TAO PyT CLI 为 NvDINOv2 提供 export、inference、train 和 default_specs。请勿启动或宣传独立的 nvdinov2 distill 操作。

未找到 AutoML 指标：TAO 的状态 KPI 会将最终训练标量报告为 train_loss。AutoML 选择时，请使用 train_loss，并将方向设为 minimize。一些 Lightning 进度行也会将同一标量显示为 train_loss_epoch；仅将其视为备用别名，而非主要监控指标。

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

Spec 参数／父模型推断

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

特定模型的交接映射：

操作	Spec 字段	推断函数	含义
export	encryption_key	key	加密密钥
export	export.checkpoint	parent_model	父任务结果文件夹中所选的 student_epoch_*.pth 检查点
export	export.onnx_file	create_onnx_file	输出 ONNX 路径
export	results_dir	output_dir	当前任务结果目录
inference	encryption_key	key	加密密钥
inference	inference.checkpoint	parent_model	父任务结果文件夹中所选的 student_epoch_*.pth 检查点
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	当前任务结果文件夹中所选的完整 model_epoch_*.pth 训练检查点

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

Deployment

• tao-deploy-nvdinov2