Tao Train Codetr

Co-DETR

Standalone install? If this session was not initialized by the TAO skill bank plugin, run the tao-setup skill first (host preflight, credentials, cross-skill discovery).

Co-DETR trains a DETR detector alongside auxiliary one-to-many assignment heads (num_co_heads), which supervise the encoder more densely than DETR's one-to-one matching alone. The auxiliary heads exist only during training; inference runs the primary DETR head. Backbones are large by default — vit_large_codetr for training, swin_large_patch4_window7_224 for the reference inference/eval configs — so this is an accuracy-first model, not a latency-first one.

Availability — probe, then fall back to the module

The codetr console script is not registered in the TAO PyTorch images checked so far (7.0.1-pyt and the 2026.7.31-rc-12-multiarch nightly), even though the module itself ships. So a bare codetr invocation fails while the network is perfectly usable.

Probe in two steps and use whichever works:

# 1. console script (preferred when present)
docker run --rm "$TAO_PYT_IMAGE" codetr --help >/dev/null 2>&1 && CODETR="codetr"

# 2. module fallback — works whenever the package is installed
[ -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; }

Both forms take identical arguments — the module entrypoint reports itself as codetr and exposes the same {train, evaluate, inference, default_specs} subtasks. Substitute $CODETR wherever this document writes codetr.

Only if both probes fail does this image genuinely lack Co-DETR. Then stop and ask which image to use — do not silently substitute another detector.

Dataclass Schemas

Generated TAO Core schemas are not yet packaged for this model, so schemas/<action>.schema.json and references/spec_template_<action>.yaml are absent except for the hand-written references/spec_template_inference.yaml. Read spec keys from the upstream experiment specs in nvidia_tao_pytorch/cv/codetr/experiment_specs/ (train.yaml, eval.yaml, inference.yaml, export.yaml) until a maintainer regenerates them. AutoML is therefore not runnable for this model.

Train Action Policy

This model is not AutoML-enabled (automl_enabled: false in references/skill_info.yaml). Run train directly; do not route through tao-skill-bank:tao-run-automl. Never add an automl_policy or workflow: key to the spec — TAO's Hydra ExperimentConfig schema rejects unknown top-level keys at config-merge time.

Supported Actions

The packaged Co-DETR CLI exposes train, evaluate, and inference. This skill exposes all three. export and TensorRT deploy flows are not available for this model.

Every action follows the standard TAO form, where anything after the spec is a Hydra override:

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

results_dir auto-appends the action name: passing results_dir=X writes to X/train/, X/evaluate/, or X/inference/. Never append the subdirectory yourself.

Training Requirements

• Dataset type: object_detection
• Formats: coco (train/evaluate), image directory + classmap (inference)
• Monitoring metric: mAP50

Per-Action Dataset Requirements

dataset.num_classes is required for every action and must match both the classmap length and the checkpoint head.

Typical Spec Overrides

# $CODETR is either `codetr` or the module form — see ## 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 is a starting point.

Getting Weights

Two fields control where training starts:

Prefer train.pretrained_model_path with the COCO detector below. It carries a trained backbone and trained detection heads, so it is the better starting point for fine-tuning on a new class set, and it is the checkpoint this skill is verified against.

A backbone-only checkpoint must match the architecture exactly. vit_large_codetr is ViT-L/16 — readable from the checkpoint as backbone.patch_embed.proj.weight (1024, 3, 16, 16), a 16x16 patch embedding. A ViT-L/14 checkpoint such as nvidia/tao/pretrained_dinov2_classification_imagenet:vit_large_patch14_dinov2 has a (1024, 3, 14, 14) embedding and cannot load into it; verify the patch size of any candidate against the target backbone before using it. For the Swin, FAN, ResNet and EfficientViT backbones, use their own matching checkpoints.

Left null, the backbone initialises randomly and needs substantially longer training.

Inference needs a full detector checkpoint, not a backbone — inference.checkpoint wants trained detection heads, which a backbone does not have. Training emits model_epoch_<N>.pth into train.results_dir.

A COCO-trained detector checkpoint

No Co-DETR detector appears under nvidia/tao on NGC. The upstream authors publish one on HuggingFace, and it loads into this build directly:

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

zongzhuofan/co-detr-vit-large-coco — Apache-2.0, a single 2.8 GB pytorch_model.pth. Verified against 7.0.1-pyt: it is the ViT-Large COCO-80 model, and pairs with

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

Loading reports 66 missing, 0 unexpected — the missing keys are the training-only collaborative heads and are expected. Use the COCO-80 classmap with it, and fold to your own classes with inference.category_mapping (above) rather than post-processing the labels.

The classmap file

A ready-made COCO-80 classmap ships with this skill at references/coco80_classmap.txt — use it directly for any COCO-trained checkpoint. It is extracted from the container's own METAINFO['classes'] in nvidia_tao_pytorch/cv/deformable_detr/model/post_process.py, which is the source of truth; do not retype a COCO list from memory, since the VOC-style names (aeroplane, motorbike) and the 91-id-with-gaps ordering both look plausible and both mis-map every class silently.

dataset.infer_data_sources.classmap is a plain-text file, one class name per line, in category_id order starting at 1 (the first foreground class). Its length must equal the number of foreground classes the model predicts — 80 for a COCO-trained checkpoint with dataset.contiguous_labels: True. These names are what inference.color_map and inference.category_mapping refer to.

Deriving the spec from a checkpoint

model.* must match the checkpoint or loading fails. Read the architecture out of the checkpoint rather than guessing — every field below is recoverable:

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

A correct load reports 0 unexpected:

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

Missing keys are normal — the collaborative heads are training-time only. Unexpected keys are not: they mean the spec describes a different architecture than the checkpoint holds.

Verified inference spec — ViT-Large, COCO-80

Confirmed against a 2.8 GB vit_large_codetr COCO checkpoint on 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"]

The published docs are ahead of the 7.0.1 image

Two fields in the online Co-DETR spec example are not in this build's schema and fail the Hydra merge with Key '<name>' not in '<Config>':

Co-DETR reuses DINO's dataset and augmentation configs, so enumerate the real fields before trusting an example:

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 hard-requires dataset.augmentation.fixed_random_crop — the ViT backbone needs a fixed input size, and build_nn_model.py raises without it. It is not optional at inference despite the name suggesting a training-time augmentation.

Important Parameters

Class indexing. The reference Co-DETR checkpoint uses 0-indexed class labels — TAO's inference path sets start_from_one=False. Verify the offset before mapping predicted class names onto another schema's integer ids.

Inference Output

With results_dir=X:

Label lines are KITTI-style — 15 fields plus a trailing confidence, boxes absolute 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>