Tao Train Rtdetr

RT-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).

RT-DETR (Real-Time DEtection TRansformer) for 2D object detection. Designed for real-time inference with competitive accuracy. Supports distillation and quantization for deployment optimization.

Set model.pretrained_backbone_path for backbone weights or train.pretrained_model_path for full model.

For TAO Deploy TensorRT actions (gen_trt_engine, TensorRT evaluate, and TensorRT inference), read references/tao-deploy-rtdetr.md first. Deploy spec templates live in this skill's references/ folder with the spec_template_deploy_*.yaml prefix.

Dataclass Schemas

Generated TAO Core schemas are packaged in schemas/<action>.schema.json, with schemas/manifest.json listing available actions. Each generated schema also emits references/spec_template_<action>.yaml from the schema top-level default field. AutoML enablement is declared at the model layer in references/skill_info.yaml via automl_enabled. Runnable AutoML for an action requires schemas/<action>.schema.json and references/spec_template_<action>.yaml to exist and parse. Use the packaged selected-action schema for automl_default_parameters, automl_disabled_parameters, defaults, min/max bounds, enums, option weights, math conditions, dependencies, and popular parameters. Do not expect ~/tao-core at runtime; maintainers regenerate schemas/templates before packaging the skill bank.

Train Action Policy

This model is AutoML-enabled at the model layer. Before handling any train-stage request, read references/skill_info.yaml and resolve the run override from either an explicit automl_policy value or the user's workflow request. Use automl_policy: on by default and only expose on / off in new launch prompts. Treat phrases like "turn off AutoML", "disable AutoML", "no HPO", or "plain training" as automl_policy: off for this run only. When automl_policy: on, automl_enabled: true, and both schemas/train.schema.json and references/spec_template_train.yaml are packaged, route the train action through tao-skill-bank:tao-run-automl by default with this model's skill_dir. Preserve workflow/application overrides for datasets, specs, output directories, GPU/platform settings, parent checkpoints, and automl_policy. Use direct model training only when automl_policy: off or the packaged train schema/template is missing; in the missing-schema case, report that AutoML is enabled but not runnable for this model until schemas are generated.

Non-train actions such as evaluate, inference, export, and deploy flows stay in this model skill. The per-run automl_policy override does not change model metadata.

Supported Actions

The packaged RT-DETR PyT CLI supports train, distill, quantize, evaluate, export, inference, and default_specs. This model skill exposes train, distill, quantize, evaluate, export, and inference; resume/retrain is performed through train with train.resume_training_checkpoint_path.

The parent PyT CLI does not expose gen_trt_engine. Use models/rtdetr/deploy for TensorRT engine generation, TensorRT evaluation, and TensorRT inference.

Training Requirements

• Dataset type: object_detection
• Formats: coco, coco_raw
• AutoML training metrics: val_mAP50 for quick operational checks or val_mAP for COCO/paper-style benchmark comparisons; both are maximized. These are the exact structured train-status KPI names.
• Standalone evaluation metrics: test_mAP50 and test_mAP. Do not use the evaluator's test_* names for AutoML trial ranking.

Per-Action Dataset Requirements

Typical Spec Overrides

Data source overrides are mandatory for every action — the agent MUST construct data source paths from the Per-Action Dataset Requirements table above and include them in spec_overrides.

S3_TRAIN = "s3://bucket/data/train"
S3_EVAL = "s3://bucket/data/eval"
CHECKPOINT = "/results/{train_job_id}/results_dir/model_epoch_000.pth"
ONNX_FILE = "/results/{export_job_id}/results_dir/rtdetr.onnx"

train (mandatory data sources):

{
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "train.gpu_ids": [0],
    "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"},
}

resume train (mandatory checkpoint):

{
    "train.num_epochs": 11,
    "train.resume_training_checkpoint_path": CHECKPOINT,
    "dataset.num_classes": "<num_classes> + 1",
    "dataset.eval_class_ids": [1, 2, 3, 4],
    "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 (mandatory data sources and checkpoint):

{
    "dataset.num_classes": "<num_classes> + 1",
    "dataset.eval_class_ids": [1, 2, 3, 4],
    "dataset.test_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
    "evaluate.checkpoint": CHECKPOINT,
}

export (mandatory checkpoint and output):

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

quantize (mandatory data sources):

{
    "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"},
    "quantize.model_path": CHECKPOINT,
}

inference (mandatory data sources and checkpoint):

{
    "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"},
    "inference.checkpoint": CHECKPOINT,
}

distill (mandatory data sources and teacher checkpoint):

{
    "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"},
    "distill.pretrained_teacher_model_path": CHECKPOINT,
}

Eval Dataset

Optional. Provides validation mAP at each checkpoint if supplied.

Important Parameters

• dataset.num_classes: Number of classes. Default 80 (MSCOCO 80-class). Must match your dataset annotations.
• model.backbone: Default resnet_50. Supported: ResNet variants, ConvNeXt, FAN, EfficientViT. RT-DETR is optimized for real-time with lighter backbones.
• train.optim.lr: Learning rate. Default 1e-4 (lower than DINO's 2e-4). lr_backbone defaults to 1e-5.
• dataset.augmentation.train_spatial_size: Training input size. Default [640, 640]. Smaller than DINO's multi-scale (up to 1333). Key to RT-DETR's speed.
• model.num_feature_levels: Default 3 (vs DINO's 4). return_interm_indices is [1,2,3].
• train.enable_ema: Exponential moving average. Default False. Enable for potentially smoother convergence.
• dataset.remap_mscoco_category: Default False. Set True only for original MSCOCO dataset with 91-to-80 category ID remapping.

Multi-GPU / Multi-Node

Launch method: torchrun (LIGHTNING_EXCLUDED_NETWORK). The entrypoint runs torchrun --nnodes=N --nproc-per-node=M train.py, NOT plain python.

• When increasing train.num_gpus, also set train.gpu_ids to the same visible device range. For example, an 8-GPU single-node Slurm run must include both "train.num_gpus": 8 and "train.gpu_ids": [0, 1, 2, 3, 4, 5, 6, 7].
• CUDA_VISIBLE_DEVICES is explicitly set (unlike Lightning-managed models which use TAO_VISIBLE_DEVICES)
• ddp with activation checkpointing: find_unused_parameters=False
• ddp without: find_unused_parameters=True
• fsdp supported, forces FP16

Multi-node env vars (set by orchestrator):

CRITICAL: NODE_RANK is copied to RANK if RANK is unset. This is required for torchrun multinode.

Export / TRT Defaults

• Export input: 640x640, opset 17
• TRT data types: FP32, FP16, INT8
• TRT workspace: 1024 MB
• TRT max_batch_size: 4

Distillation

RT-DETR supports knowledge distillation with a teacher model. Requires distill action with distill.pretrained_teacher_model_path and a distillation binding configuration.

Use the packaged references/spec_template_distill.yaml as the starting point. The validated default binding uses the RT-DETR distiller's explicit IOU feature path:

distill:
  bindings:
  - student_module_name: srcs
    teacher_module_name: srcs
    criterion: IOU
    weight: 1.0

Do not substitute DINO-style output names such as pred_logits / pred_boxes, and do not bind arbitrary decoder heads unless you have verified the module returns captured feature lists. The RT-DETR distiller asserts that IOU bindings must use srcs or dsrcs.