tao-train-segformer

Par nvidia · skills

SegFormer pour la segmentation sémantique. Architecture légère basée sur les transformers avec des caractéristiques hiérarchiques

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

SegFormer

SegFormer pour la segmentation sémantique. Architecture légère basée sur les transformers avec extraction hiérarchique des features. Efficace pour les tâches de segmentation en temps réel.

Définir model.backbone.pretrained_backbone_path pour les poids du backbone.

Pour les actions TAO Deploy TensorRT (gen_trt_engine, évaluation TensorRT et inférence TensorRT), lisez d'abord references/tao-deploy-segformer.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.

Schémas Dataclass

Les schémas TAO Core générés sont packagés dans schemas/<action>.schema.json, avec schemas/manifest.json qui liste les actions disponibles. Chaque schéma généré émet aussi references/spec_template_<action>.yaml à partir du champ default de 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'AutoML exécutable nécessite toujours que schemas/train.schema.json et references/spec_template_train.yaml existent et se chargent. Utilisez le schéma train packagé pour automl_default_parameters, automl_disabled_parameters, les valeurs par défaut, les limites min/max, les énumérations, les poids d'option, les conditions mathématiques, les dépendances et les paramètres populaires. Ne comptez pas sur ~/tao-core au moment de l'exécution ; les mainteneurs régénèrent les schémas/templates avant de packager la skill bank.

Politique d'Action Train

Ce modèle est activé AutoML au niveau du modèle. Avant de traiter toute demande de stage train, lisez references/skill_info.yaml et résolvez le run override à partir d'une valeur automl_policy explicite ou de la demande de workflow de l'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 schemas/train.schema.json et references/spec_template_train.yaml sont packagé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 overrides workflow/application pour les datasets, les specs, les répertoires de sortie, les paramètres GPU/platform, les checkpoints parents et automl_policy. Utilisez l'entraînement direct du modèle uniquement quand automl_policy: off ou quand le schéma/template train packagé est manquant ; 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 modèle. L'override automl_policy par exécution ne change pas les métadonnées du modèle.

Exigences d'Entraînement

  • Type de dataset : segmentation
  • Formats : unet
  • Métrique de monitoring : val_miou

Exigences de Dataset par Action

Action Clé Spec Source Fichiers Liste ?
evaluate dataset.segment.root_dir eval_dataset Non
export dataset.segment.root_dir train_datasets Non
inference dataset.segment.root_dir eval_dataset Non
quantize dataset.segment.root_dir train_datasets Non
quantize dataset.segment.quant_calibration_dataset.images_dir train_datasets Non
train dataset.segment.root_dir train_datasets Non

Overrides de Spec Typiques

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 (sources de données obligatoires) :

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "dataset.segment.batch_size": 4,
    "dataset.segment.root_dir": f"{S3_TRAIN}",
}

evaluate (sources de données obligatoires) :

{
    "evaluate.batch_size": 4,
    "dataset.segment.root_dir": f"{S3_EVAL}",
}

gen_trt_engine :

{
    "gen_trt_engine.tensorrt.data_type": "fp16",
}

inference (sources de données obligatoires) :

{
    "dataset.segment.batch_size": 1,
    "dataset.segment.root_dir": f"{S3_EVAL}",
}

export (sources de données obligatoires) :

{
    "dataset.segment.root_dir": f"{S3_TRAIN}",
}

quantize (sources de données obligatoires) :

{
    "dataset.segment.root_dir": f"{S3_TRAIN}",
    "dataset.segment.quant_calibration_dataset.images_dir": f"{S3_TRAIN}",
}

Dataset d'Évaluation

Optionnel. Les données de validation font généralement partie de la structure root_dir.

Paramètres Importants

  • dataset.segment.num_classes : Nombre de classes de segmentation. Par défaut 2 (binaire). Doit correspondre au nombre de classes dans vos annotations de masque.
  • model.backbone.type : Par défaut fan_small_12_p4_hybrid. Les variantes supportées incluent FAN variants, SegFormer MIT variants et autres.
  • dataset.segment.root_dir : Répertoire racine du dataset de segmentation.
  • dataset.segment.img_size : Taille de l'image d'entrée. Par défaut 256. Augmentez pour une segmentation plus fine au prix de la mémoire.
  • train.optim.lr : Taux d'apprentissage. Par défaut 6e-5.
  • model.freeze_backbone : Désactiver ou non le backbone lors de l'entraînement. Utile pour le fine-tuning avec des données limitées.
  • dataset.segment.batch_size : Taille du batch par GPU. Par défaut 8.

Multi-GPU / Multi-Node

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

Clé Spec Description Défaut
train.num_gpus Nombre de GPUs 1
train.gpu_ids Indices de dispositif GPU [0]
train.num_nodes Nombre de nœuds 1
train.sync_batchnorm Sync BN sur les GPUs configurable
train.use_distributed_sampler Utiliser un distributed sampler configurable
  • Stratégie multi-GPU : ddp_find_unused_parameters_true
  • Pas de support fsdp

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

Matériel

Minimum 1 GPU(s), recommandé 2 GPU(s). 16 GB+ (V100 ou A100) VRAM par GPU. SegFormer est relativement léger. L'img_size par défaut=256 est économe en mémoire. Augmentez img_size pour une résolution plus élevée au prix de la mémoire et de la vitesse.

Patterns d'Erreurs

CUDA out of memory : Réduisez batch_size ou img_size. La mémoire SegFormer évolue quadratiquement avec la taille de l'image.

Décalage num_classes : Assurez-vous que dataset.segment.num_classes correspond au nombre réel de classes dans vos annotations de masque.

Inférence des Paramètres 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 microservices infer_params.py.

Mappings d'inférence depuis TAO Core segformer.config.json :

Action Champ Spec Fonction d'Inférence Signification
evaluate encryption_key key clé de chiffrement
evaluate evaluate.checkpoint parent_model fichier modèle déduit du dossier de résultats du job parent
evaluate evaluate.trt_engine parent_model fichier modèle déduit 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 modèle déduit 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 modèle déduit 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 modèle déduit du dossier de résultats du job parent
inference inference.trt_engine parent_model fichier modèle déduit du dossier de résultats du job parent
inference results_dir output_dir répertoire de résultats du job actuel
quantize encryption_key key clé de chiffrement
quantize quantize.model_path parent_model fichier modèle déduit du dossier de résultats du job parent
quantize results_dir output_dir répertoire de résultats du job actuel
train encryption_key key clé de chiffrement
train model.backbone.pretrained_backbone_path ptm_if_no_resume_model PTM quand aucun checkpoint de reprise n'existe
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 modèle déduit 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 modèle sélectionné. N'ajoutez pas ces mappings de retour à config.json et ne patcherez pas les scripts runners générés pour deviner les chemins de checkpoint.

Déploiement

  • tao-deploy-segformer — Workflow de déploiement SegFormer pour la génération du moteur TensorRT, l'évaluation TensorRT et l'inférence TensorRT avec TAO Deploy.

Skills similaires