tao-train-mask-auto-encoder

Par nvidia · skills

Masked Auto-Encoder (MAE) pour le pré-entraînement auto-supervisé et le fine-tuning. Masque des patches aléatoires et reconstruit

npx skills add https://github.com/nvidia/skills --skill tao-train-mask-auto-encoder

MAE

MAE (Masked Autoencoder) pour préentraînement et affinage auto-supervisés. Masque des patches aléatoires et les reconstruit pour apprendre des représentations visuelles. Supporte les étapes pretrain et finetune.

Définissez train.pretrained_model_path pour les poids MAE préentraînés lors de l'affinage.

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

Schémas Dataclass

Les schémas TAO Core générés sont empaquetés dans schemas/<action>.schema.json, avec schemas/manifest.json listant les actions disponibles. Chaque schéma généré émet également references/spec_template_<action>.yaml à partir du champ default du niveau supérieur du schéma. L'activation AutoML est déclarée au niveau du modèle dans references/skill_info.yaml via automl_enabled. L'exécution runnable d'AutoML nécessite toujours que schemas/train.schema.json et references/spec_template_train.yaml existent et se parsent. Utilisez le schéma train empaqueté pour automl_default_parameters, automl_disabled_parameters, les valeurs par défaut, les limites min/max, les enums, les poids d'option, les conditions mathématiques, les dépendances et les paramètres populaires. Ne s'attendez pas à ~/tao-core à l'exécution ; les mainteneurs régénèrent les schémas/modèles avant d'empaqueter la banque de skills.

Politique d'Action Train

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

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

Exigences d'Entraînement

  • Type de dataset : image_classification
  • Formats : ssl
  • Intentions de dataset acceptées : training, evaluation, testing
  • Métrique de monitoring : train_loss

Exigences de Dataset par Action

Action Clé Spec Source Fichiers Liste ?
train dataset.train_data_sources train_datasets images_train.tar.gz Non
train dataset.val_data_sources eval_dataset images_val.tar.gz Non
evaluate dataset.val_data_sources eval_dataset images_val.tar.gz Non
inference dataset.test_data_sources inference_dataset images_test.tar.gz Non

Substitutions Spec Typiques

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

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

train (sources de données obligatoires) :

{
    "dataset.train_data_sources": f"{S3_TRAIN}/images_train.tar.gz",
    "dataset.val_data_sources": f"{S3_EVAL}/images_val.tar.gz",
    "train.num_epochs": 10,
    "train.optim.lr": 2e-4,
}

evaluate (sources de données obligatoires) :

{
    "dataset.val_data_sources": f"{S3_EVAL}/images_val.tar.gz",
}

inference (sources de données obligatoires) :

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

Dataset d'Évaluation

Optionnel. Le préentraînement n'a pas besoin de données d'évaluation. L'affinage utilise optionnellement l'ensemble de validation.

Paramètres Importants

  • train.stage: Étape d'entraînement. Options : pretrain, finetune. Pretrain apprend les représentations via masquage. Finetune ajoute une tête de classification.
  • model.arch: Architecture. Par défaut convnextv2_base. Large gamme d'options incluant ConvNeXt, Hiera, variantes ViT.
  • model.num_classes: Nombre de classes pour l'affinage. Par défaut 1000 (ImageNet). Pertinent uniquement en étape finetune.
  • model.mask_ratio: Fraction de patches à masquer lors du préentraînement. Typiquement 0.75.
  • model.norm_pix_loss: Normaliser ou non les valeurs de pixels dans la perte de reconstruction.
  • train.optim.lr: Taux d'apprentissage. Par défaut 2e-4.
  • dataset.augmentation: Paramètres d'augmentation incluant mixup, cutmix pour l'affinage.

Multi-GPU / Multi-Nœud

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

Clé Spec Description Par défaut
train.num_gpus Nombre de GPUs 1
train.gpu_ids Indices des devices GPU [0]
train.num_nodes Nombre de nœuds 1
train.distributed_strategy ddp ou fsdp ddp
  • ddp utilise find_unused_parameters=True
  • fsdp force FP16
  • Multi-GPU fortement recommandé pour le préentraînement (grandes tailles de batch nécessaires)

Variables d'env multi-nœud (définies par l'orchestrateur) : WORLD_SIZE, NODE_RANK, MASTER_ADDR, MASTER_PORT, NUM_GPU_PER_NODE.

Matériel

Minimum 2 GPU, recommandé 8 GPU. 24 GB+ (A100 recommandé) de VRAM par GPU. Le préentraînement MAE bénéficie de grandes tailles de batch sur plusieurs GPUs. L'affinage demande des ressources plus modestes.

Motifs d'Erreurs

Inadéquation d'étape : Assurez-vous que train.stage correspond à votre intention (pretrain vs finetune). L'affinage sans pretrained_model_path entraîne à partir de zéro.

Inadéquation num_classes (finetune uniquement) : Assurez-vous que model.num_classes correspond au nombre de classes de votre dataset lors de l'affinage.

Inférence de Param Spec / Modèle Parent

Les mappings 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 mappings avec les helpers SDK avant create_job(). Cela reflète l'ancien flux infer_params.py des microservices.

Mappings d'inférence à partir de mae.config.json TAO Core :

Action Champ Spec Fonction d'Inférence Signification
evaluate encryption_key key clé de chiffrement
evaluate evaluate.checkpoint parent_model fichier de modèle inféré à partir du dossier de résultats du job parent
evaluate evaluate.trt_engine parent_model fichier de modèle inféré à partir du dossier de résultats du job parent
evaluate results_dir output_dir répertoire de résultats du job actuel
export encryption_key key clé de chiffrement
export export.checkpoint parent_model fichier de modèle inféré à partir du dossier de résultats du job parent
export export.onnx_file create_onnx_file chemin ONNX de sortie
export results_dir output_dir répertoire de résultats du job actuel
gen_trt_engine encryption_key key clé de chiffrement
gen_trt_engine gen_trt_engine.onnx_file parent_model fichier de modèle inféré à partir du dossier de résultats du job parent
gen_trt_engine gen_trt_engine.trt_engine create_engine_file chemin du moteur TensorRT de sortie
gen_trt_engine results_dir output_dir répertoire de résultats du job actuel
inference encryption_key key clé de chiffrement
inference inference.checkpoint parent_model fichier de modèle inféré à partir du dossier de résultats du job parent
inference inference.trt_engine parent_model fichier de modèle inféré à partir du dossier de résultats du job parent
inference results_dir output_dir répertoire de résultats du job actuel
train encryption_key key clé de chiffrement
train results_dir output_dir répertoire de résultats du job actuel
train train.pretrained_model_path ptm_if_no_resume_model PTM quand aucun checkpoint de reprise n'existe
train train.resume_training_checkpoint_path resume_model fichier de modèle inféré à partir du dossier de résultats du job actuel

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 mappings à config.json et ne patcherez pas les scripts de runner générés pour deviner les chemins de checkpoint.

Déploiement

Skills similaires