
名称	tao-train-nvdinov2
描述	NVDINOv2用于自监督视觉表示学习。通过自蒸馏（教师-学生）无需标签地训练视觉Transformer，并生成通用视觉特征。当需要为TAO NVDINOv2骨干网络进行训练、导出或运行推理时使用。触发短语包括"train NVDINOv2", “self-supervised ViT pretraining”, “DINOv2 backbone”, “visual representation learning”。
开源协议	Apache-2.0 compatibility: 需要docker + nvidia-container-toolkit。 metadata:
版本	“0.1.0”
作者	NVIDIA Corporation allowed-tools: Read Bash tags: - self - supervised - learning

NVDINOv2

独立安装？ 如果此会话不是由TAO技能库插件初始化的，请先运行tao-setup技能（主机预检、凭据、跨技能发现）。

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

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

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

数据类模式

生成的TAO Core模式打包在schemas/<action>.schema.json中，schemas/manifest.json列出了可用的操作。每个生成的模式还会从模式顶层的default字段发出references/spec_template_<action>.yaml。AutoML启用状态在模型层的references/skill_info.yaml中通过automl_enabled声明。操作可运行的AutoML要求schemas/<action>.schema.json和references/spec_template_<action>.yaml存在并可解析。对于automl_default_parameters、automl_disabled_parameters、默认值、最小/最大边界、枚举、选项权重、数学条件、依赖项和流行参数，请使用打包的所选操作模式。不要在运行时期望~/tao-core；维护人员在打包技能库之前会重新生成模式/模板。

训练操作策略

此模型在模型层启用了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或打包的训练模式/模板缺失时，才使用直接模型训练；在缺失模式的情况下，报告AutoML已启用但在此模型生成模式之前不可运行。

非训练操作（如inference、export和部署流程）保留在此模型技能中。每次运行的automl_policy覆盖不会改变模型元数据。

训练要求

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

按操作的数据集要求

操作	规范键	来源	来源产物 → 运行时值	列表?
推理	dataset.test_dataset.images_dir	推理数据集	`images_test.tar.gz` → 提取的图像文件夹	否
训练	dataset.train_dataset.images_dir	训练数据集	`images_train.tar.gz` → 提取的图像文件夹	否

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

典型规范覆盖

数据源覆盖对于训练和推理是强制性的——代理必须根据上面的按操作数据集要求表构造数据源路径，并将它们包含在spec_overrides中。

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

训练（强制数据源）：

{
    "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.checkpoint": "<selected train/AutoML student_epoch_* checkpoint>",
    "export.onnx_file": "/path/to/results/nvdinov2.onnx",
}

推理（强制数据源）：

{
    "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>",
    "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生成工作进程）。

规范键	描述	默认值
`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训练非常消耗内存。需要A100 40GB以上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。最新的文件是训练检查点，推理加载器会报告意外的键，例如state_dict、优化器状态和调度器状态。

导出检查点具有意外的Lightning键：导出也使用所选的student_epoch_*.pth检查点。完整的model_epoch_*.pth检查点仅用于通过train.resume_training_checkpoint_path恢复/重新训练。

将TensorRT引擎传递给PyT推理：打包的PyT nvdinov2 inference实现仅加载.pth或.tlt模型路径。TAO Deploy gen_trt_engine为下游消费者构建TensorRT引擎，但PyT推理操作不会在该引擎上运行。

没有单独的蒸馏操作：当前TAO PyT CLI为NvDINOv2公开了export、inference、train和default_specs。不要启动或宣传独立的nvdinov2 distill操作。

未找到AutoML指标：TAO的状态KPI将最终训练标量报告为train_loss。使用train_loss并采用最小化方向进行AutoML选择。一些Lightning进度行也会将相同的标量渲染为train_loss_epoch；仅将其视为回退别名，而不是主要监控指标。

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

规范参数 / 父模型推断

模型特定的推断映射属于此MD文件，而不在config.json中。生成的运行器应在create_job()之前阅读本节并与SDK助手应用映射。这类似于旧的微服务infer_params.py流程。

模型特定的交接映射：

操作	规范字段	推断函数	含义
导出	`encryption_key`	`key`	加密密钥
导出	`export.checkpoint`	`parent_model`	来自父作业结果文件夹的所选`student_epoch_*.pth`检查点
导出	`export.onnx_file`	`create_onnx_file`	输出ONNX路径
导出	`results_dir`	`output_dir`	当前作业结果目录
推理	`encryption_key`	`key`	加密密钥
推理	`inference.checkpoint`	`parent_model`	来自父作业结果文件夹的所选`student_epoch_*.pth`检查点
推理	`results_dir`	`output_dir`	当前作业结果目录
训练	`encryption_key`	`key`	加密密钥
训练	`results_dir`	`output_dir`	当前作业结果目录
训练	`train.pretrained_model_path`	`ptm_if_no_resume_model`	当没有恢复检查点时的PTM
训练	`train.resume_training_checkpoint_path`	`resume_model`	来自当前作业结果文件夹的所选完整`model_epoch_*.pth`训练检查点

对于parent_model或parent_model_folder，将上游训练/导出/AutoML子作业ID作为parent_job_id传递。SDK列出父结果文件夹，过滤检查点工件，并返回所选模型文件或文件夹。不要将这些映射添加回config.json，也不要修补生成的运行器脚本以猜测检查点路径。

部署

tao-deploy-nvdinov2