tao-train-nvdinov2

Par nvidia · skills

NVDINOv2 pour l'apprentissage de représentations visuelles auto-supervisé. Entraîne des transformers de vision par auto-distillation

npx skills add https://github.com/nvidia/skills --skill tao-train-nvdinov2

NVDINOv2

NVDINOv2 pour l'apprentissage de représentations visuelles auto-supervisé. Entraîne des vision transformers via l'auto-distillation (maître-étudiant) sans labels. Produit des features visuelles à usage général.

Définissez train.pretrained_model_path pour les poids ViT pré-entraînés.

Pour les actions TAO Deploy TensorRT (gen_trt_engine), consultez d'abord references/tao-deploy-nvdinov2.md. Les templates de spec de déploiement se trouvent dans le dossier references/ de cette skill avec le préfixe spec_template_deploy_*.yaml.

Dataclass Schemas

Les schemas TAO Core générés sont packagés dans schemas/<action>.schema.json, avec schemas/manifest.json listant les actions disponibles. Chaque schema généré émet aussi references/spec_template_<action>.yaml à partir du champ default au niveau supérieur du schema. L'activation AutoML est déclarée au niveau du modèle dans references/skill_info.yaml via automl_enabled. L'exécution d'AutoML requiert toujours que schemas/train.schema.json et references/spec_template_train.yaml existent et se parsent. Utilisez le schema train packagé pour automl_default_parameters, automl_disabled_parameters, les valeurs par défaut, les bornes min/max, les énumérations, les poids d'option, les conditions mathématiques, les dépendances et les paramètres populaires. N'attendez pas ~/tao-core à l'exécution ; les mainteneurs régénèrent les schemas/templates avant de packager la skill bank.

Train Action Policy

Ce modèle est activé pour AutoML au niveau du modèle. Avant de traiter une demande au stade train, lisez references/skill_info.yaml et résolvez l'override d'exécution à partir soit d'une valeur automl_policy explicite, soit de la demande de workflow utilisateur. Traitez les phrases comme « turn off AutoML », « disable AutoML », « no HPO » ou « plain training » comme automl_policy: off pour cette exécution uniquement ; sinon, utilisez par défaut auto. Quand automl_policy: auto, automl_enabled: true, et que les deux schemas/train.schema.json et references/spec_template_train.yaml sont packagés, routez l'action train par défaut via tao-skill-bank:tao-run-automl avec le skill_dir de ce modèle. Préservez les overrides workflow/application pour les datasets, specs, répertoires de sortie, paramètres GPU/plateforme, checkpoints parents et automl_policy. Utilisez l'entraînement direct du modèle uniquement quand automl_policy: off ou que le schema train/template packagé est absent ; dans le cas du schema manquant, signalez qu'AutoML est activé mais non exécutable pour ce modèle jusqu'à ce que les schemas soient générés.

Les actions non-train telles que evaluate, inference, export et les flux de déploiement restent dans cette skill de modèle. L'override automl_policy par exécution ne change pas les métadonnées du modèle.

Training Requirements

  • Dataset type: image_classification
  • Formats: ssl
  • Monitoring metric: train_loss

Per-Action Dataset Requirements

Action Spec Key Source Files List?
distill dataset.train_dataset.images_dir train_datasets images_train.tar.gz No
inference dataset.test_dataset.images_dir inference_dataset images_test.tar.gz No
train dataset.train_dataset.images_dir train_datasets images_train.tar.gz No

Typical Spec Overrides

Les overrides de source de données sont obligatoires pour chaque action — l'agent DOIT construire les chemins de source de données à partir du tableau Per-Action Dataset Requirements ci-dessus et les inclure dans spec_overrides.

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

train (mandatory data sources):

{
    "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 (mandatory data sources):

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

inference (mandatory data sources):

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

Eval Dataset

Optionnel. L'entraînement SSL n'utilise pas de labels. L'évaluation est spécifique à la tâche en aval.

Important Parameters

  • model.backbone.teacher_type: Variante ViT maître. Par défaut vit_l (ViT-Large).
  • model.backbone.student_type: Variante ViT étudiant. Par défaut vit_l. Correspond généralement au maître.
  • model.backbone.img_size: Taille d'image d'entrée. Par défaut 518. Une résolution plus élevée produit de meilleures features mais coûte plus de mémoire.
  • model.backbone.patch_size: Taille de patch ViT. Par défaut 14.
  • dataset.batch_size: Taille de batch par GPU. Par défaut 4. L'entraînement SSL est intensif en mémoire en raison des passes forward duales (maître+étudiant).
  • train.layerwise_decay: Décroissance du taux d'apprentissage par couche. Important pour le fine-tuning ViT.
  • train.clip_grad_norm: Clipping du gradient. Important pour un entraînement SSL stable.

Multi-GPU / Multi-Node

Méthode de lancement: Lightning-managed (processus python unique, Lightning crée les workers).

Spec Key Description Default
train.num_gpus Number of GPUs 1
train.gpu_ids GPU device indices [0]
train.num_nodes Number of nodes 1
  • Strategy: auto (Lightning choisit automatiquement la meilleure stratégie)
  • sync_batchnorm est toujours activé — critique pour l'entraînement SSL avec framework maître-étudiant
  • Multi-GPU fortement recommandé (4-8 GPUs) pour un entraînement SSL significatif

Multi-node env vars (défini par l'orchestrateur): WORLD_SIZE, NODE_RANK, MASTER_ADDR, MASTER_PORT, NUM_GPU_PER_NODE.

Hardware

Minimum 4 GPU(s), recommandé 8 GPU(s). 40GB+ (A100 recommandé) VRAM par GPU. SSL avec maître+étudiant ViT-Large est très intensif en mémoire. Nécessite des GPUs A100 40GB+. Multi-GPU fortement recommandé.

Error Patterns

CUDA out of memory: Maître+étudiant ViT-Large avec img_size=518 requiert 40GB+ de mémoire GPU. Réduisez batch_size, img_size, ou utilisez une variante ViT plus petite.

Slow convergence: SSL a besoin de nombreuses epochs. La valeur par défaut 10 est pour des tests rapides ; les exécutions en production utilisent typiquement 100+ epochs.

Spec Param / Parent Model Inference

Les mappages d'inférence spécifiques au modèle appartiennent à ce fichier MD, pas à config.json. Les runners générés doivent lire cette section et appliquer les mappages avec les helpers SDK avant create_job(). Cela reflète l'ancien flux microservices infer_params.py.

Mappages d'inférence depuis TAO Core nvdinov2.config.json:

Action Spec Field Inference Function Meaning
distill encryption_key key encryption key
distill model.distill.pretrained_non_distill_pl_model_path parent_model model file inferred from the parent job results folder
distill results_dir output_dir current job results directory
export encryption_key key encryption key
export export.checkpoint parent_model model file inferred from the parent job results folder
export export.onnx_file create_onnx_file output ONNX path
export results_dir output_dir current job results directory
inference encryption_key key encryption key
inference inference.checkpoint parent_model model file inferred from the parent job results folder
inference inference.trt_engine parent_model model file inferred from the parent job results folder
inference results_dir output_dir current job results directory
train encryption_key key encryption key
train results_dir output_dir current job results directory
train train.pretrained_model_path ptm_if_no_resume_model PTM when no resume checkpoint exists
train train.resume_training_checkpoint_path resume_model model file inferred from the current job results folder

Pour parent_model ou parent_model_folder, passez l'id du job enfant train/export/AutoML en amont comme parent_job_id. Le SDK liste le dossier de résultats parent, filtre les artefacts de checkpoint, et retourne le fichier ou dossier de modèle sélectionné. N'ajoutez pas ces mappages à config.json et ne patchchez pas les scripts runner générés pour deviner les chemins de checkpoint.

Deployment

  • tao-deploy-nvdinov2 — Workflow de déploiement NvDINOv2 pour la génération de moteur TensorRT avec TAO Deploy.

Skills similaires