tao-train-codetr

Par nvidia · skills

Co-DETR (CoDINO) pour la détection d'objets. Un détecteur de la famille DETR avec hybridation collaborative

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

Co-DETR

Installation autonome ? Si cette session n'a pas été initialisée par le plugin skill bank TAO, exécutez d'abord le skill tao-setup (vérification des prérequis de l'hôte, identifiants, découverte inter-skills).

Co-DETR entraîne un détecteur DETR aux côtés de têtes d'assignation auxiliaires un-vers-plusieurs (num_co_heads), qui supervisent l'encoder de manière plus dense que le seul appariement un-vers-un de DETR. Les têtes auxiliaires n'existent que pendant l'entraînement ; l'inférence exécute la tête DETR principale. Les backbones sont volumineux par défaut — vit_large_codetr pour l'entraînement, swin_large_patch4_window7_224 pour les configs d'inférence/évaluation de référence — il s'agit donc d'un modèle orienté précision, non latence.

Disponibilité — vérifier, puis basculer vers le module

Le console script codetr n'est pas enregistré dans les images TAO PyTorch vérifiées jusqu'à présent (7.0.1-pyt et la nightly 2026.7.31-rc-12-multiarch), bien que le module lui-même soit livré. Une invocation codetr nue échoue tandis que le réseau est parfaitement utilisable.

Vérifiez en deux étapes et utilisez celle qui fonctionne :

# 1. console script (préféré quand présent)
docker run --rm "$TAO_PYT_IMAGE" codetr --help >/dev/null 2>&1 && CODETR="codetr"

# 2. fallback module — fonctionne dès que le package est installé
[ -z "${CODETR:-}" ] && docker run --rm --entrypoint sh "$TAO_PYT_IMAGE" \
  -c 'python3 -c "import nvidia_tao_pytorch.cv.codetr"' >/dev/null 2>&1 \
  && CODETR="python3 -m nvidia_tao_pytorch.cv.codetr.entrypoint.codetr"

[ -z "${CODETR:-}" ] && { echo "FATAL: Co-DETR not available in $TAO_PYT_IMAGE"; exit 1; }

Les deux formes acceptent des arguments identiques — le point d'entrée du module se rapporte à codetr et expose les mêmes sous-tâches {train, evaluate, inference, default_specs}. Substituez $CODETR partout où ce document écrit codetr.

Ce n'est que si les deux vérifications échouent que l'image manque véritablement de Co-DETR. Arrêtez-vous et demandez quelle image utiliser — ne remplacez pas silencieusement par un autre détecteur.

Schémas Dataclass

Les schémas TAO Core générés ne sont pas encore empaquetés pour ce modèle, donc schemas/<action>.schema.json et references/spec_template_<action>.yaml sont absents sauf pour le manuel references/spec_template_inference.yaml. Lisez les clés spec à partir des specs d'expérience en amont dans nvidia_tao_pytorch/cv/codetr/experiment_specs/ (train.yaml, eval.yaml, inference.yaml, export.yaml) jusqu'à ce qu'un mainteneur les régénère. AutoML n'est donc pas exécutable pour ce modèle.

Politique de l'action Train

Ce modèle n'est pas activé pour AutoML (automl_enabled: false dans references/skill_info.yaml). Exécutez l'entraînement directement ; ne routez pas via tao-skill-bank:tao-run-automl. N'ajoutez jamais de clé automl_policy ou workflow: au spec — le schéma Hydra ExperimentConfig de TAO rejette les clés top-level inconnues lors de la fusion de config.

Actions Supportées

Le CLI Co-DETR empaqueté expose train, evaluate et inference. Ce skill expose les trois. export et les flux de déploiement TensorRT ne sont pas disponibles pour ce modèle.

Chaque action suit la forme TAO standard, où tout après le spec est une override Hydra :

$CODETR <action> -e /absolute/path/spec.yaml [key=value ...]

results_dir ajoute automatiquement le nom de l'action : passer results_dir=X écrit dans X/train/, X/evaluate/ ou X/inference/. N'ajoutez jamais le sous-répertoire vous-même.

Exigences d'Entraînement

  • Type de dataset : object_detection
  • Formats : coco (train/evaluate), répertoire d'images + classmap (inference)
  • Métrique de suivi : mAP50

Exigences de Dataset par Action

Action Clé Spec Fichiers
train dataset.train_data_sources image_dir, json_file (COCO)
train dataset.val_data_sources image_dir, json_file (COCO)
evaluate dataset.test_data_sources image_dir, json_file (COCO)
inference dataset.infer_data_sources.image_dir répertoire d'images
inference dataset.infer_data_sources.classmap liste de classes séparée par des retours à la ligne, un nom par ligne

dataset.num_classes est requis pour chaque action et doit correspondre à la fois à la longueur du classmap et à la tête du checkpoint.

Overrides Spec Typiques

# $CODETR est soit `codetr` soit la forme module — voir ## Availability
$CODETR inference -e "$SPEC" \
  inference.checkpoint=/abs/codetr.pth \
  dataset.infer_data_sources.image_dir=/abs/images \
  dataset.infer_data_sources.classmap=/abs/classmap.txt \
  results_dir=/abs/results \
  inference.conf_threshold=0.3 \
  inference.num_gpus=8

references/spec_template_inference.yaml est un point de départ.

Obtenir les Poids

Deux champs contrôlent d'où l'entraînement commence :

Champ Contient
model.pretrained_backbone_path un checkpoint backbone-only initialisé model.backbone
train.pretrained_model_path un checkpoint complet Co-DETR (ou compatible DETR-family) à affiner

Préférez train.pretrained_model_path avec le détecteur COCO ci-dessous. Il porte un backbone entraîné et des têtes de détection entraînées, c'est donc un meilleur point de départ pour l'affinage sur un nouvel ensemble de classes, et c'est le checkpoint contre lequel ce skill est vérifié.

Un checkpoint backbone-only doit correspondre exactement à l'architecture. vit_large_codetr est ViT-L/16 — lisible depuis le checkpoint sous la forme backbone.patch_embed.proj.weight (1024, 3, 16, 16), un embedding de patch 16x16. Un checkpoint ViT-L/14 tel que nvidia/tao/pretrained_dinov2_classification_imagenet:vit_large_patch14_dinov2 a un embedding (1024, 3, 14, 14) et ne peut pas charger dedans ; vérifiez la taille de patch de tout candidat par rapport au backbone cible avant de l'utiliser. Pour les backbones Swin, FAN, ResNet et EfficientViT, utilisez leurs checkpoints correspondants.

Laissé null, le backbone initialise aléatoirement et nécessite un entraînement substantiellement plus long.

L'inférence a besoin d'un checkpoint détecteur complet, pas seulement un backboneinference.checkpoint veut des têtes de détection entraînées, qu'un backbone n'a pas. L'entraînement émet model_epoch_<N>.pth dans train.results_dir.

Un checkpoint détecteur entraîné sur COCO

Aucun détecteur Co-DETR n'apparaît sous nvidia/tao sur NGC. Les auteurs en amont en publient un sur HuggingFace, et il charge directement dans ce build :

huggingface-cli download zongzhuofan/co-detr-vit-large-coco pytorch_model.pth --local-dir <dir>

zongzhuofan/co-detr-vit-large-coco — Apache-2.0, un unique pytorch_model.pth de 2,8 GB. Vérifié contre 7.0.1-pyt : c'est le modèle ViT-Large COCO-80, et s'associe avec

model:
  backbone: vit_large_codetr
  num_queries: 1500
  num_feature_levels: 5
  return_interm_indices: [0, 1, 2, 3, 4]
  num_co_heads: 1
dataset:
  num_classes: 80

Le chargement rapporte 66 missing, 0 unexpected — les clés manquantes sont les têtes collaboratives training-only et sont attendues. Utilisez le classmap COCO-80 avec celui-ci, et repliez vers vos propres classes avec inference.category_mapping (ci-dessus) plutôt que post-traiter les labels.

Le fichier classmap

Un classmap COCO-80 prêt à l'emploi est livré avec ce skill à references/coco80_classmap.txt — utilisez-le directement pour tout checkpoint entraîné sur COCO. Il est extrait du conteneur lui-même METAINFO['classes'] dans nvidia_tao_pytorch/cv/deformable_detr/model/post_process.py, qui est la source de vérité ; ne retapez pas une liste COCO de mémoire, puisque les noms style VOC (aeroplane, motorbike) et l'ordre 91-id-avec-gaps semblent tous deux plausibles et remappent silencieusement chaque classe.

dataset.infer_data_sources.classmap est un fichier texte brut, un nom de classe par ligne, dans l'ordre category_id commençant à 1 (la première classe de premier plan). Sa longueur doit égaler le nombre de classes de premier plan que le modèle prédit — 80 pour un checkpoint entraîné sur COCO avec dataset.contiguous_labels: True. Ces noms sont ceux auxquels font référence inference.color_map et inference.category_mapping.

Dériver le spec depuis un checkpoint

model.* doit correspondre au checkpoint ou le chargement échoue. Lisez l'architecture depuis le checkpoint plutôt que de deviner — chaque champ ci-dessous est récupérable :

sd = torch.load(ckpt, map_location="cpu", weights_only=False)["state_dict"]
# backbone.patch_embed.proj.weight (1024,3,16,16) -> ViT-L/16      -> vit_large_codetr
# query_head.transformer.query_embed.weight (1500,256)             -> num_queries 1500
# max index in query_head...{encoder,decoder}.layers.N.            -> enc_layers / dec_layers
# query_head...cls_branches.0.weight (80,256)                      -> dataset.num_classes 80
# roi_head.<i>. indices present                                    -> num_co_heads (count)
# neck.p2..p6                                                      -> num_feature_levels 5

Un chargement correct rapporte 0 unexpected :

Missing keys: 66 total (66 expected for collab heads / buffers, 0 unexpected)

Les clés manquantes sont normales — les têtes collaboratives sont training-time only. Les clés unexpected ne le sont pas : elles signifient que le spec décrit une architecture différente du checkpoint.

Spec d'inférence vérifiée — ViT-Large, COCO-80

Confirmé contre un checkpoint COCO vit_large_codetr de 2,8 GB sur 7.0.1-pyt :

results_dir: /abs/results
model:
  backbone: vit_large_codetr
  num_queries: 1500
  num_feature_levels: 5
  return_interm_indices: [0, 1, 2, 3, 4]
  two_stage_type: standard
  num_co_heads: 1
  hidden_dim: 256
  nheads: 8
  enc_layers: 6
  dec_layers: 6
  dim_feedforward: 2048
dataset:
  num_classes: 80
  batch_size: 2
  augmentation:
    fixed_padding: true
    fixed_random_crop: 1536        # REQUIRED by vit_large_codetr — see below
    input_mean: [0.485, 0.456, 0.406]
    input_std: [0.229, 0.224, 0.225]
    test_random_resize: 1280
  infer_data_sources:
    image_dir: /abs/images
    classmap: /abs/coco_classmap.txt
inference:
  checkpoint: /abs/codetr_pytorch_model.pth
  conf_threshold: 0.3
  input_width: 640
  input_height: 640
  num_gpus: 2
  category_mapping:
    car: ["car", "bus", "truck"]

La documentation publiée est en avance sur l'image 7.0.1

Deux champs dans l'exemple spec Co-DETR en ligne ne sont pas dans le schéma de ce build et échouent la fusion Hydra avec Key '<name>' not in '<Config>' :

Champ Statut dans 7.0.1-pyt
dataset.contiguous_labels absent de DINODatasetConfig
dataset.augmentation.pad_size_divisor absent de DINOAugmentationConfig

Co-DETR réutilise les configs dataset et augmentation de DINO, énumérez donc les champs réels avant de faire confiance à un exemple :

python3 -c "
from nvidia_tao_pytorch.config.dino.dataset import DINODatasetConfig, DINOAugmentationConfig
import dataclasses
print([f.name for f in dataclasses.fields(DINODatasetConfig)])
print([f.name for f in dataclasses.fields(DINOAugmentationConfig)])"

vit_large_codetr requiert obligatoirement dataset.augmentation.fixed_random_crop — le backbone ViT a besoin d'une taille d'entrée fixe, et build_nn_model.py lève une exception sans cela. C'est obligatoire à l'inférence malgré le nom suggérant une augmentation training-time.

Paramètres Importants

Paramètre Notes
model.backbone vit_large_codetr (défaut train) ou swin_large_patch4_window7_224 (défaut inference/eval). Doit correspondre au checkpoint.
model.num_queries 1500 pour la config ViT, 900 pour Swin. Associé au backbone — ne mélangez pas.
model.num_feature_levels / return_interm_indices 5 / [0,1,2,3,4] pour ViT ; 4 / [1,2,3,4] pour Swin.
model.num_co_heads Têtes collaboratives auxiliaires. 2 pour la config ViT train, 1 pour la config inference de référence. Training-time only.
model.soft_nms_enabled La config train active soft-NMS (linear, IoU 0,8) pour correspondre à la config test Co-DETR original.
inference.conf_threshold Les détections au-dessous de ce seuil sont supprimées au moment de l'écriture. Défaut TAO 0,5.
dataset.num_classes Doit correspondre à la longueur du classmap et à la tête du checkpoint.

Indexation des classes. Le checkpoint Co-DETR de référence utilise des labels de classe 0-indexés — le chemin d'inférence de TAO définit start_from_one=False. Vérifiez le décalage avant de mapper les noms de classe prédits sur les ids entiers d'un autre schéma.

Sortie d'Inférence

Avec results_dir=X :

Artefact Localisation
Labels de détection KITTI X/inference/labels/<image_stem>.txt
Images annotées X/inference/images_annotated/
Config d'exécution + statut X/inference/experiment.yaml, X/inference/status.json

Les lignes de label sont de style KITTI — 15 champs plus une confiance finale, boîtes absolues xyxy :

<class_name> 0.00 0 0.00 <x1> <y1> <x2> <y2> 0.00 0.00 0.00 0.00 0.00 0.00 0.00 <score>

Vocabulaire Prédit

Co-DETR prédit le vocabulaire sur lequel son checkpoint a été entraîné — COCO-80 pour le checkpoint de référence, dans l'ordre donné par dataset.infer_data_sources.classmap. Les labels émis portent les noms de classe, pas les indices.

Replier vers un ensemble de classes plus petit — utilisez inference.category_mapping

Quand un modèle en aval s'entraîne sur un ensemble plus petit, repliez ici, non après :

inference:
  category_mapping:
    bicycle: ["bicycle", "motorcycle"]
    car:     ["car", "bus", "truck"]
    person:  ["person"]

Les valeurs sont des listes. Sémantique, depuis model/category_mapping.py :

Comportement Détail
originaux non mappés supprimés
un nom dans deux groupes garde le premier, enregistre un avertissement
un nom absent du classmap avertit, ignore — la correspondance est exacte, y compris la casse
un remap vide lève ValueError
ids de catégorie de sortie 0..K-1 dans l'ordre où le mapping est écrit

Ensuite elle exécute apply_category_mapping_groupnms — soft-NMS par catégorie de sortie après la fusion. C'est pourquoi replier ici vaut mieux que renommer les labels après : un objet détecté à la fois comme truck et car devient deux boîtes de la même classe l'instant où elles se replioient ensemble, et seule une NMS post-repli supprime le doublon. Un changement de nom en aval expédie les doublons plus loin.

Ce dédoublonnage a besoin de model.soft_nms_enabled: True, qui n'est pas le défaut. Le schéma le livre False, et avec celui-ci éteint le repli renomme les classes sans rien fusionner — le même résultat qu'un changement de nom en aval donnerait, doublons inclus. Replier {vehicle: ["car", "bus", "truck"]} sur 8 images de trafic à conf_threshold: 0.05 :

model.soft_nms_enabled boîtes émises
False (défaut schéma) 277 — exactement 146 car + 102 truck + 29 bus, rien fusionné
True 151 — 126 doublons supprimés

Activez-le quand category_mapping groupe des classes que le détecteur confond entre elles, ce qui pour les véhicules COCO le fait de manière fiable. soft_nms_iou_threshold (défaut 0,8) contrôle l'agressivité de la fusion.

Les clés inference.color_map et les noms de classe dans les labels émis suivent tous les deux les catégories de sortie une fois ceci activé.

Un fichier de label vide est significatif et est toujours écrit : en KITTI cela signifie « cette image n'a pas d'objets ».

Convertir les labels pour l'entraînement en aval

KITTI n'est rarement le format final. Les deux conversions sont publiées dans le conteneur des services de données TAO via annotations convert -e <spec> :

  • KITTI → COCOdata.input_format: KITTI, data.output_format: COCO, plus kitti.{image_dir, label_dir, mapping}. kitti.mapping est une liste YAML de dicts à clé unique dont les valeurs sont des listes- car: [car], jamais - car: car. Le convertisseur construit sa table de recherche avec {label: k for k, v in cat_map.items() for label in v}, en itérant la valeur, donc une chaîne brute produit les noms de classe c, a, r : rien ne correspond, chaque boîte est supprimée, et l'exécution imprime toujours Execution status: PASS et quitte 0.
  • COCO → ODVG (entraînement Grounding DINO) — data.input_format: COCO, data.output_format: ODVG, plus coco.ann_file. Émet <name>_odvg.jsonl et <name>_odvg_labelmap.json.

Il n'y a pas de conversion directe KITTI → ODVG ; chainez les deux.

Multi-GPU / Multi-Node

Définissez train.num_gpus / inference.num_gpus. gpu_spec_key est train.num_gpus. Les grands backbones rendent ce modèle gourmand en mémoire — réduisez dataset.batch_size avant de réduire le nombre de GPU en cas de OOM.

Matériel

Orienté précision avec de grands backbones et des compteurs de requête élevés (900–1500). Attendez-vous à une mémoire et une latence substantiellement plus élevées que DINO ou RT-DETR à taille d'entrée égale. train.activation_checkpoint échange du calcul contre de la mémoire quand l'entraînement dépasse la capacité.

Motifs d'Erreur

Symptôme Cause Correction
codetr: command not found Console script non enregistré — attendu sur les builds courants ; le module est généralement toujours livré Utilisez python3 -m nvidia_tao_pytorch.cv.codetr.entrypoint.codetr. Seulement si l'import échoue aussi Co-DETR est véritablement absent
inference/inference/ imbriqué Avez ajouté le nom de l'action à results_dir Passez seulement le répertoire parent
Fichiers de label vides partout conf_threshold au-dessus de la plage de score du modèle Baissez-le (les pipelines de référence utilisent 0,3)
Noms de classe décalés d'un Le checkpoint de référence est 0-indexé Vérifiez le décalage lors du mapping vers un autre schéma
Mismatch taille/forme lors du chargement du checkpoint model.backbone ne correspond pas au checkpoint Alignez backbone, num_queries, num_feature_levels, return_interm_indices, num_co_heads
Key 'X' not in 'DINODatasetConfig' / 'DINOAugmentationConfig' Spec utilise un champ documenté absent de ce build Énumérez les vrais champs dataclass ; contiguous_labels et pad_size_divisor sont documentés mais pas dans 7.0.1
vit_large_codetr requires dataset.augmentation.fixed_random_crop Le backbone ViT a besoin d'une taille d'entrée fixe Définissez fixed_random_crop (1536 pour la config ViT de référence) — requis à l'inférence aussi
Clés unexpected non-zéro au chargement L'architecture spec diffère du checkpoint Re-dérivez model.* depuis les tenseurs du checkpoint ; les clés manquantes seules sont normales
CUDA OOM à l'entraînement Grand backbone + compteur de requête élevé Baissez dataset.batch_size ; activez train.activation_checkpoint
Chemins non trouvés dans le conteneur Mismatch chemin hôte/conteneur Montez pour que les chemins soient identiques ; confirmez que les images sont sous le montage, pas seulement le spec

Déploiement

export et le déploiement TensorRT ne sont pas disponibles pour Co-DETR dans ce build. Pour un détecteur déployable, entraînez Co-DETR comme teacher et distillez dans un student qui supporte export — tao-train-rtdetr et tao-train-dino exposent tous les deux export et gen_trt_engine.

Références

  • references/checkpoint-spec-pairing.md — dériver backbone, num_queries, num_feature_levels et num_classes depuis les tenseurs d'un checkpoint. Une discordance échoue silencieusement : PASS, exit 0, labels vides.

Skills similaires