tao-train-rtdetr

Par nvidia · skills

RT-DETR (Real-Time DEtection TRansformer) pour la détection d'objets 2D. Conçu pour l'inférence en temps réel avec

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

RT-DETR

RT-DETR (Real-Time DEtection TRansformer) pour la détection d'objets 2D. Conçu pour l'inférence en temps réel avec une précision compétitive. Supporte la distillation et la quantification pour l'optimisation du déploiement.

Définissez model.pretrained_backbone_path pour les poids du backbone ou train.pretrained_model_path pour le modèle complet.

Pour les actions TAO Deploy TensorRT (gen_trt_engine, évaluation TensorRT et inférence TensorRT), lisez d'abord references/tao-deploy-rtdetr.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 packagé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 au niveau supérieur du schéma. L'activation d'AutoML est déclarée au niveau du modèle dans references/skill_info.yaml via automl_enabled. 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 d'entraînement packagé pour automl_default_parameters, automl_disabled_parameters, les valeurs par défaut, les limites min/max, les énumérations, les poids d'options, 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 schémas/modèles avant de packager la skill bank.

Politique d'action Train

Ce modèle est activé pour AutoML au niveau du modèle. Avant de traiter toute requête de stage d'entraînement, lisez references/skill_info.yaml et résolvez l'override d'exécution à partir soit d'une valeur explicite automl_policy, soit de la requête 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, utilisez la valeur par défaut auto. Lorsque automl_policy: auto, automl_enabled: true et que schemas/train.schema.json et references/spec_template_train.yaml sont packagés, routez l'action d'entraînement 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, 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 lorsque automl_policy: off ou que le schéma/modèle d'entraînement 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 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.

Conditions d'entraînement

  • Type de dataset : object_detection
  • Formats : coco, coco_raw
  • Métrique de monitoring : val_mAP50

Conditions de dataset par action

Action Clé de spec Source Fichiers Liste ?
distill dataset.train_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations.json Oui
distill dataset.val_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json Non
evaluate dataset.test_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json Non
gen_trt_engine gen_trt_engine.tensorrt.calibration.cal_image_dir calibration_dataset images.tar.gz Oui
inference dataset.infer_data_sources inference_dataset image_dir: images.tar.gz, classmap: label_map.txt Non
quantize dataset.train_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations.json Oui
quantize dataset.val_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json Non
quantize dataset.quant_calibration_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations.json Non
train dataset.train_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations.json Oui
train dataset.val_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json Non

Overrides de spec typiques

Les overrides de sources de données sont obligatoires pour chaque action — l'agent DOIT construire les chemins des sources de données à partir du tableau Conditions 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) :

{
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.num_classes": "<num_classes> + 1",
    "dataset.train_data_sources": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
    "dataset.val_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

evaluate (sources de données obligatoires) :

{
    "dataset.num_classes": "<num_classes> + 1",
    "dataset.test_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

export :

{
    "dataset.num_classes": "<num_classes> + 1",
    "export.input_height": 640,
    "export.input_width": 640,
}

quantize (sources de données obligatoires) :

{
    "dataset.num_classes": "<num_classes> + 1",
    "quantize.layers": [
        {
            "module_name": "*",
            "weights": {
                "dtype": "float8_e4m3fn"
            },
            "activations": {
                "dtype": "float8_e4m3fn"
            }
        }
    ],
    "dataset.train_data_sources": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
    "dataset.val_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
    "dataset.quant_calibration_data_sources": {"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"},
}

gen_trt_engine (sources de données obligatoires) :

{
    "gen_trt_engine.tensorrt.data_type": "FP16",
    "gen_trt_engine.tensorrt.calibration.cal_image_dir": [f"{S3_TRAIN}/images.tar.gz"],
}

inference (sources de données obligatoires) :

{
    "dataset.num_classes": "<num_classes> + 1",
    "dataset.infer_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "classmap": f"{S3_EVAL}/label_map.txt"},
}

distill (sources de données obligatoires) :

{
    "dataset.train_data_sources": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
    "dataset.val_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

Dataset d'évaluation

Optionnel. Fournit la mAP de validation à chaque checkpoint s'il est fourni.

Paramètres importants

  • dataset.num_classes : Nombre de classes. Par défaut 80 (MSCOCO 80 classes). Doit correspondre à vos annotations de dataset.
  • model.backbone : Par défaut resnet_50. Supportés : variantes ResNet, ConvNeXt, FAN, EfficientViT. RT-DETR est optimisé pour le temps réel avec des backbones plus légers.
  • train.optim.lr : Taux d'apprentissage. Par défaut 1e-4 (plus bas que le 2e-4 de DINO). lr_backbone par défaut à 1e-5.
  • dataset.augmentation.train_spatial_size : Taille d'entrée d'entraînement. Par défaut [640, 640]. Plus petit que le multi-échelle de DINO (jusqu'à 1333). Clé de la vitesse de RT-DETR.
  • model.num_feature_levels : Par défaut 3 (vs 4 pour DINO). return_interm_indices est [1,2,3].
  • train.enable_ema : Moyenne mobile exponentielle. Par défaut False. Activez pour une convergence potentiellement plus lisse.
  • dataset.remap_mscoco_category : Par défaut False. Définissez à True uniquement pour le dataset MSCOCO original avec remappage des ID de catégorie de 91 à 80.

Multi-GPU / Multi-nœud

Méthode de lancement : torchrun (LIGHTNING_EXCLUDED_NETWORK). Le point d'entrée exécute torchrun --nnodes=N --nproc-per-node=M train.py, PAS le simple python.

Clé de spec Description Par défaut
train.num_gpus Nombre de GPUs par nœud 1
train.gpu_ids Indices des appareils GPU [0]
train.num_nodes Nombre de nœuds 1
train.distributed_strategy ddp ou fsdp ddp
  • CUDA_VISIBLE_DEVICES est explicitement défini (contrairement aux modèles gérés par Lightning qui utilisent TAO_VISIBLE_DEVICES)
  • ddp avec activation checkpointing : find_unused_parameters=False
  • ddp sans : find_unused_parameters=True
  • fsdp supporté, force FP16

Variables d'environnement multi-nœud (définies par l'orchestrateur) :

Variable Objectif
WORLD_SIZE Nombre de nœuds (déclenche le mode multinœud)
NODE_RANK Rang de ce nœud (indexé à partir de 0)
MASTER_ADDR IP du nœud rang-0
MASTER_PORT Port rang-0 (par défaut 29500)
NUM_GPU_PER_NODE GPUs par nœud (par défaut : tous les visibles)

CRITIQUE : NODE_RANK est copié vers RANK si RANK n'est pas défini. Ceci est requis pour torchrun multinœud.

Export / Valeurs par défaut TRT

  • Export d'entrée : 640x640, opset 17
  • Types de données TRT : FP32, FP16, INT8
  • Espace de travail TRT : 1024 MB
  • max_batch_size TRT : 4

Référence TAO Deploy complète : tao-deploy-rtdetr.

Distillation

RT-DETR supporte la distillation de connaissances avec un modèle enseignant. Nécessite l'action distill avec le chemin du modèle enseignant et la configuration des liaisons de distillation.

Matériel

Minimum 1 GPU, recommandé 2 GPU. 16GB+ (V100 ou A100) VRAM par GPU. RT-DETR est plus efficace en mémoire que DINO/GDINO en raison de la taille d'entrée plus petite (640x640) et de moins de niveaux de features. S'entraîne bien sur un seul GPU pour les datasets de petite à moyenne taille.

Motifs d'erreur

CUDA out of memory : Réduisez batch_size. RT-DETR à 640x640 est plus léger que DINO à 1333px, mais batch_size > 8 peut toujours causer une OOM sur les GPUs 16GB.

Incompatibilité num_classes : RT-DETR utilise par défaut 80 (pas 91 comme DINO). Assurez-vous que dataset.num_classes correspond à vos catégories d'annotations.

return_interm_indices vs num_feature_levels : Par défaut [1,2,3] avec num_feature_levels=3. Doit être cohérent si modifié.

Inférence de spec param / modèle parent

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(). Ceci reflète le flux ancien infer_params.py des microservices.

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

Action Champ de spec Fonction d'inférence Signification
distill distill.pretrained_teacher_model_path parent_model fichier de modèle déduit du dossier des résultats du job parent
distill encryption_key key clé de chiffrement
distill results_dir output_dir répertoire des résultats du job courant
evaluate encryption_key key clé de chiffrement
evaluate evaluate.checkpoint parent_model fichier de modèle déduit du dossier des résultats du job parent
evaluate evaluate.trt_engine parent_model fichier de modèle déduit du dossier des résultats du job parent
evaluate results_dir output_dir répertoire des résultats du job courant
export encryption_key key clé de chiffrement
export export.checkpoint parent_model fichier de modèle déduit du dossier des résultats du job parent
export export.onnx_file create_onnx_file chemin de sortie ONNX
export results_dir output_dir répertoire des résultats du job courant
gen_trt_engine encryption_key key clé de chiffrement
gen_trt_engine gen_trt_engine.onnx_file parent_model fichier de modèle déduit du dossier des résultats du job parent
gen_trt_engine gen_trt_engine.tensorrt.calibration.cal_cache_file create_cal_cache chemin du cache de calibration
gen_trt_engine gen_trt_engine.trt_engine create_engine_file chemin du moteur TensorRT en sortie
gen_trt_engine results_dir output_dir répertoire des résultats du job courant
inference encryption_key key clé de chiffrement
inference inference.checkpoint parent_model fichier de modèle déduit du dossier des résultats du job parent
inference inference.trt_engine parent_model fichier de modèle déduit du dossier des résultats du job parent
inference results_dir output_dir répertoire des résultats du job courant
quantize encryption_key key clé de chiffrement
quantize quantize.model_path parent_model fichier de modèle déduit du dossier des résultats du job parent
quantize results_dir output_dir répertoire des résultats du job courant
train encryption_key key clé de chiffrement
train model.pretrained_backbone_path ptm_if_no_resume_model PTM quand aucun checkpoint de reprise n'existe
train results_dir output_dir répertoire des résultats du job courant
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 déduit du dossier des résultats du job courant

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

Skills similaires