⚡

RT-DETR目标检测训练Skill tao-train-rtdetr

该技能基于 NVIDIA TAO 框架，为 RT-DETR（Real-Time DEtection TRansformer）实时目标检测模型提供从训练、评估、蒸馏、量化、导出到推理的全流程支持。RT-DETR 专为低延迟、高精度的 2D 物体检测场景设计，支持 ResNet/ConvNeXt 等多种骨干网络、多 GPU/多节点分布式训练、AutoML 自动调优以及 TensorRT 部署优化。适用于智能交通、机器人、工业质检等实时视觉应用。关键词：RT-DETR、TAO、实时目标检测、模型训练、模型评估、知识蒸馏、模型量化、TensorRT、COCO、AutoML、多GPU训练、低延迟检测。

视觉模型训练 0 次安装 0 次浏览 更新于 9/6/2026
名称 tao-train-rtdetr
描述 RT-DETR（实时检测Transformer）用于2D物体检测。专为实时推理设计，具有竞争力的精度，并支持用于部署优化的蒸馏和量化。当需要为TAO RT-DETR模型进行训练、评估、蒸馏、量化、导出或运行推理时使用。触发短语包括“train RT-DETR”、“real-time DETR”、“low-latency object detection”、“RT-DETR distillation / quantization”。
开源协议 Apache-2.0 compatibility: 需要 docker + nvidia-container-toolkit。 metadata:
版本 “0.1.0”
作者 NVIDIA Corporation allowed-tools: Read Bash tags: - object - detection

RT-DETR

独立安装？ 如果此会话不是由 TAO 技能库插件初始化的，请先运行 tao-setup 技能（主机预检、凭据、跨技能发现）。

RT-DETR（实时检测Transformer）用于2D物体检测。专为实时推理设计，具有竞争力的精度，并支持用于部署优化的蒸馏和量化。

设置 model.pretrained_backbone_path 来指定骨干网络权重，或设置 train.pretrained_model_path 来指定完整模型。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请先阅读 references/tao-deploy-rtdetr.md。部署规格模板位于此技能的 references/ 文件夹中，文件名以 spec_template_deploy_*.yaml 开头。

数据类 Schema

生成的 TAO Core schema 打包在 schemas/<action>.schema.json 中，schemas/manifest.json 列出了可用操作。每个生成的 schema 还会从 schema 顶层的 default 字段输出 references/spec_template_<action>.yaml。AutoML 的启用状态在 references/skill_info.yaml 的模型层通过 automl_enabled 声明。某个操作可运行 AutoML 的前提是：schemas/<action>.schema.json 和 references/spec_template_<action>.yaml 存在且可解析。使用打包的选定操作 schema 获取 automl_default_parameters、automl_disabled_parameters、默认值、最小/最大边界、枚举、选项权重、数学条件、依赖关系和常用参数。不要期望运行时存在 ~/tao-core；维护者在打包技能库之前会重新生成 schema/模板。

训练操作策略

本模型在模型层启用了 AutoML。在处理任何训练阶段请求之前，请先阅读 references/skill_info.yaml，并从显式的 automl_policy 值或用户的工作流请求中解析运行覆盖。默认使用 automl_policy: on，仅在新启动提示中暴露 on / off。像“关闭 AutoML”、“禁用 AutoML”、“不要 HPO”或“普通训练”这样的短语，针对当前运行应视为 automl_policy: off。当 automl_policy: on、automl_enabled: true 且 schemas/train.schema.json 与 references/spec_template_train.yaml 都已打包时，默认应通过 tao-skill-bank:tao-run-automl 路由训练操作并使用此模型的 skill_dir。保留数据集、规格、输出目录、GPU/平台设置、父检查点和 automl_policy 的工作流/应用覆盖。只有当 automl_policy: off 或打包的训练 schema/模板缺失时，才直接使用模型训练；在 schema 缺失的情况下，报告该模型的 AutoML 已启用但在生成 schema 之前无法运行。

非训练操作如 evaluate、inference、export 和部署流程仍保留在此模型技能中。每次运行的 automl_policy 覆盖不会改变模型元数据。

支持的操作

打包的 RT-DETR PyT CLI 支持 train、distill、quantize、evaluate、export、inference 和 default_specs。此模型技能暴露 train、distill、quantize、evaluate、export 和 inference；通过 train 配合 train.resume_training_checkpoint_path 实现恢复/再训练。

父 PyT CLI 不暴露 gen_trt_engine。使用 models/rtdetr/deploy 进行 TensorRT 引擎生成、TensorRT 评估和 TensorRT 推理。

训练要求

  • 数据集类型： object_detection
  • 格式： coco, coco_raw
  • 监控指标： mAP50 用于快速操作检查；val_mAP 用于 COCO/论文式基准对比。

各操作的数据集要求

操作 规格键 来源 文件 列表？
distill dataset.train_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations.json 是
distill dataset.val_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json 否
evaluate dataset.test_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json 否
inference dataset.infer_data_sources inference_dataset image_dir: images.tar.gz, classmap: label_map.txt 是
quantize dataset.train_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations.json 是
quantize dataset.val_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json 否
quantize dataset.quant_calibration_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations.json 否
train dataset.train_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations.json 是
train dataset.val_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json 否

典型规格覆盖

数据源覆盖每个操作都必须提供 —— 代理必须根据上表“各操作的数据集要求”构建数据源路径，并在 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.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"},
}

恢复训练（必需检查点）：

{
    "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"},
}

评估（必需的数据源和检查点）：

{
    "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,
}

导出（必需的检查点和输出）：

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

量化（必需的数据源）：

{
    "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,
}

推理（必需的数据源和检查点）：

{
    "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,
}

蒸馏（必需的数据源和教师检查点）：

{
    "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,
}

评估数据集

可选。如果提供，会在每个检查点输出验证 mAP。

重要参数

  • dataset.num_classes：类别数量。默认 80（MSCOCO 80 类）。必须与您的数据集标注匹配。
  • model.backbone：默认 resnet_50。支持的骨干：ResNet 变体、ConvNeXt、FAN、EfficientViT。RT-DETR 针对较轻骨干的实时性进行了优化。
  • train.optim.lr：学习率。默认 1e-4（低于 DINO 的 2e-4）。lr_backbone 默认 1e-5。
  • dataset.augmentation.train_spatial_size：训练输入大小。默认 [640, 640]。小于 DINO 的多尺度（可达 1333）。这是 RT-DETR 速度的关键。
  • model.num_feature_levels：默认 3（DINO 是 4）。return_interm_indices 为 [1,2,3]。
  • train.enable_ema：指数移动平均。默认 False。启用后可能使收敛更平滑。
  • dataset.remap_mscoco_category：默认 False。仅当用于原始 MSCOCO 数据集并需要进行 91 到 80 的类别 ID 重映射时才设为 True。

多 GPU / 多节点

启动方式： torchrun（LIGHTNING_EXCLUDED_NETWORK）。入口点运行 torchrun --nnodes=N --nproc-per-node=M train.py，而不是普通的 python。

规格键 说明 默认
train.num_gpus 每个节点的 GPU 数 1
train.gpu_ids GPU 设备索引 [0]
train.num_nodes 节点数 1
train.distributed_strategy ddp 或 fsdp ddp
  • 增加 train.num_gpus 时，同时设置 train.gpu_ids 为相同的可见设备范围。例如，单节点 8 GPU 的 Slurm 运行必须同时包含 "train.num_gpus": 8 和 "train.gpu_ids": [0, 1, 2, 3, 4, 5, 6, 7]。
  • CUDA_VISIBLE_DEVICES 会被显式设置（不同于 Lightning 管理的模型使用 TAO_VISIBLE_DEVICES）。
  • ddp 配合激活检查点：find_unused_parameters=False
  • ddp 无激活检查点：find_unused_parameters=True
  • fsdp 支持，强制 FP16

多节点环境变量（由编排器设置）：

变量 用途
WORLD_SIZE 节点数（触发多节点模式）
NODE_RANK 本节点排名（从0开始）
MASTER_ADDR 排名0的节点 IP
MASTER_PORT 排名0的端口（默认29500）
NUM_GPU_PER_NODE 每节点 GPU 数（默认：所有可见）

关键点： 如果 RANK 未设置，NODE_RANK 会被复制到 RANK。这是 torchrun 多节点所必需的。

导出 / TRT 默认值

  • 导出输入：640x640，opset 17
  • TRT 数据类型：FP32, FP16, INT8
  • TRT 工作空间：1024 MB
  • TRT 最大批大小：4

蒸馏

RT-DETR 支持使用教师模型进行知识蒸馏。需要 distill 操作并设置 distill.pretrained_teacher_model_path 和一个蒸馏绑定配置。

使用打包的 references/spec_template_distill.yaml 作为起点。经过验证的默认绑定使用 RT-DETR 蒸馏器的显式 IOU 特征路径：

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

不要替换为 DINO 风格的输出名称，如 pred_logits / pred_boxes，也不要绑定任意的解码器头，除非您已验证模块返回捕获的特征列表。RT-DETR 蒸馏器断言 IOU 绑定必须使用 srcs 或 dsrcs。

硬件

最低 1 GPU，推荐 2 GPU。每 GPU 16GB+（V100 或 A100）显存。由于输入尺寸较小（640x640）和特征层级较少，RT-DETR 比 DINO/GDINO 内存效率更高。对于中小型数据集，在单 GPU 上也能很好训练。

错误模式

CUDA 内存不足：减小 batch_size。RT-DETR 在 640x640 下比 DINO 在 1333px 下轻量，但在 16GB GPU 上 batch_size > 8 仍可能 OOM。

num_classes 不匹配：RT-DETR 默认 80（不是 DINO 的 91）。确保 dataset.num_classes 与您的标注类别匹配。

来自类别 ID 的 CUDA 索引断言：如果 COCO 类别 ID 是 1 基的或未重映射为从 0 开始的连续 ID，请将 dataset.num_classes 设置为 max(category_id) + 1，并将 dataset.eval_class_ids 与实际类别 ID 对齐。对于打包的带 ID 1-4 的四类 S3 样例，使用 dataset.num_classes: 5 和 dataset.eval_class_ids: [1, 2, 3, 4]。

return_interm_indices 与 num_feature_levels 不一致：默认是 [1,2,3] 且 num_feature_levels=3。如果更改必须一致。

导出形状不匹配：除非模型已针对不同形状进行训练和检查，否则保持 RT-DETR 导出和部署消费端输入尺寸为已验证的 640x640 默认值。较旧打包的 960x544 模板形状在 ONNX 追踪时可能因 hybrid_encoder.py 位置嵌入加法中的“张量大小不匹配”而失败。

AutoML 指标提取：RT-DETR 在结构化的训练状态和日志中输出检测指标。对于 COCO/论文式基准比较，优化 val_mAP，方向为最大化；对于明确的 AP50 工作流，优化 mAP50。在解析原始日志之前，优先使用 results_dir/train/status.json 或 AutoML 结果状态。对于默认检测模型调用，不要优化 val_loss。

检查点交接：对于 evaluate/export/inference/quantize/distill/resume，请使用检查点解析器作用于最佳 AutoML 子作业的 results_dir/train/ 文件夹，并选择适合操作的 model_epoch_*.pth 检查点。RT-DETR 也可能写入一个 latest 符号链接，但只有在调用者明确请求 latest 时才应使用。保持 dataset.num_classes、dataset.eval_class_ids、model.num_queries 和 model.num_select 与训练一致。

父级 rtdetr gen_trt_engine 被 PyT CLI 拒绝：在已验证的 7.0.0 PyT 容器中，rtdetr gen_trt_engine 不是有效的父模型子任务。使用 RT-DETR 部署工作流（references/tao-deploy-rtdetr.md）进行 TensorRT 引擎生成、TensorRT 评估和 TensorRT 推理。

规格参数 / 父模型推断

模型特定的推断映射应位于此 MD 文件中，而不是 config.json。生成的运行器应在 create_job() 之前阅读此部分，并使用 SDK 帮助器应用映射。这类似于旧微服务 infer_params.py 流程。

来自 TAO Core rtdetr.config.json 的推断映射：

操作 规格字段 推断函数 含义
distill distill.pretrained_teacher_model_path parent_model 从父作业结果文件夹推断出的模型文件
distill encryption_key key 加密密钥
distill results_dir output_dir 当前作业结果目录
evaluate encryption_key key 加密密钥
evaluate evaluate.checkpoint parent_model 从父作业结果文件夹推断出的模型文件
evaluate evaluate.trt_engine parent_model 从父作业结果文件夹推断出的模型文件
evaluate results_dir output_dir 当前作业结果目录
export encryption_key key 加密密钥
export export.checkpoint parent_model 从父作业结果文件夹推断出的模型文件
export export.onnx_file create_onnx_file 输出 ONNX 路径
export results_dir output_dir 当前作业结果目录
inference encryption_key key 加密密钥
inference inference.checkpoint parent_model 从父作业结果文件夹推断出的模型文件
inference inference.trt_engine parent_model 从父作业结果文件夹推断出的模型文件
inference results_dir output_dir 当前作业结果目录
quantize encryption_key key 加密密钥
quantize quantize.model_path parent_model 从父作业结果文件夹推断出的模型文件
quantize results_dir output_dir 当前作业结果目录
train encryption_key key 加密密钥
train model.pretrained_backbone_path ptm_if_no_resume_model 当无恢复检查点时的 PTM
train results_dir output_dir 当前作业结果目录
train train.pretrained_model_path ptm_if_no_resume_model 当无恢复检查点时的 PTM
train train.resume_training_checkpoint_path resume_model 从当前作业结果文件夹推断出的模型文件

对于 parent_model 或 parent_model_folder，将上游训练/导出/AutoML 子作业 ID 作为 parent_job_id 传递。SDK 列出父结果文件夹，过滤检查点工件，并返回选定的模型文件或文件夹。不要将这些映射添加回 config.json，也不要修补生成的运行器脚本以猜测检查点路径。

部署