tao-train-rtdetr

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。

数据类架构

生成的 TAO Core 架构打包在 schemas/<action>.schema.json 中，schemas/manifest.json 列出了可用操作。每个生成的架构还会根据架构顶层的 default 字段生成 references/spec_template_<action>.yaml。AutoML 是否启用由模型层的 references/skill_info.yaml 中的 automl_enabled 声明。某项操作要能运行 AutoML，必须存在并且可解析 schemas/<action>.schema.json 和 references/spec_template_<action>.yaml。请使用所选操作的已打包架构，获取 automl_default_parameters、automl_disabled_parameters、默认值、最小值／最大值边界、枚举、选项权重、数学条件、依赖项和常用参数。不要假设运行时存在 ~/tao-core；维护者会在打包技能库之前重新生成架构／模板。

训练操作策略

此模型在模型层启用了 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 或缺少已打包的训练架构／模板时，才直接训练模型；若缺少架构／模板，请说明 AutoML 已启用，但在为此模型生成架构之前无法运行。

evaluate、inference、export 和部署流程等非训练操作仍由此模型技能处理。每次运行的 automl_policy 覆盖设置不会更改模型元数据。

支持的操作

已打包的 RT-DETR PyT CLI 支持 train、distill、quantize、evaluate、export、inference 和 default_specs。此模型技能提供 train、distill、quantize、evaluate、export 和 inference；恢复／重新训练通过设置 train.resume_training_checkpoint_path 使用 train 完成。

上级 PyT CLI 不提供 gen_trt_engine。请使用 models/rtdetr/deploy 生成 TensorRT 引擎、执行 TensorRT 评估和 TensorRT 推理。

训练要求

• 数据集类型： object_detection
• 格式： coco、coco_raw
• AutoML 训练指标： 快速运行检查使用 val_mAP50，COCO／论文风格的基准比较使用 val_mAP；两者均以最大化为目标。这些是结构化训练状态 KPI 的准确名称。
• 独立评估指标： test_mAP50 和 test_mAP。请勿使用评估器的 test_* 名称对 AutoML 试验进行排名。

各操作的数据集要求

操作	Spec 键	来源	文件	是否为列表？
蒸馏	dataset.train_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations.json	是
蒸馏	dataset.val_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	否
评估	dataset.test_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	否
推理	dataset.infer_data_sources	inference_dataset	image_dir: images.tar.gz, classmap: label_map.txt	是
量化	dataset.train_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations.json	是
量化	dataset.val_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	否
量化	dataset.quant_calibration_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations.json	否
训练	dataset.train_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations.json	是
训练	dataset.val_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	否

常见 Spec 覆盖项

每项操作都必须覆盖数据源——Agent 必须根据上方的各操作数据集要求表构造数据源路径，并将其包含在 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 数据集且需要将类别 ID 从 91 映射到 80 时，才设为 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（与使用 TAO_VISIBLE_DEVICES 的 Lightning 管理模型不同）
• 使用激活检查点时的 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 数量（默认值：所有可见 GPU）

重要： 如果未设置 RANK，则会将 NODE_RANK 复制到 RANK。这是 torchrun 多节点运行所必需的。

导出 / TRT 默认值

• 导出输入：640x640，opset 17
• TRT 数据类型：FP32、FP16、INT8
• TRT 工作区：1024 MB
• TRT max_batch_size：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。640x640 下的 RT-DETR 比 1333px 下的 DINO 更轻量，但在 16GB GPU 上，batch_size > 8 仍可能导致内存溢出。

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 对齐。对于附带的四类 S3 示例（ID 为 1-4），请使用 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 的位置编码相加处报错：The size of tensor a (...) must match the size of tensor b (...)。

AutoML 指标提取：RT-DETR 会在结构化训练状态和日志中输出检测指标。进行 COCO / 论文风格的基准比较时，请使用 direction: maximize 优化 val_mAP；对于明确采用 AP50 的工作流，请优化 val_mAP50。独立评估会输出对应的 test_mAP 和 test_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 与训练时保持一致。

PyT CLI 拒绝父级 rtdetr gen_trt_engine 命令：在经过验证的 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.pretrained_teacher_model_path	parent_model	从父作业结果文件夹推断出的模型文件
蒸馏	encryption_key	key	加密密钥
蒸馏	results_dir	output_dir	当前作业结果目录
评估	encryption_key	key	加密密钥
评估	evaluate.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
评估	evaluate.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
评估	results_dir	output_dir	当前作业结果目录
导出	encryption_key	key	加密密钥
导出	export.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
导出	export.onnx_file	create_onnx_file	ONNX 输出路径
导出	results_dir	output_dir	当前作业结果目录
推理	encryption_key	key	加密密钥
推理	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
推理	inference.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
推理	results_dir	output_dir	当前作业结果目录
量化	encryption_key	key	加密密钥
量化	quantize.model_path	parent_model	从父作业结果文件夹推断出的模型文件
量化	results_dir	output_dir	当前作业结果目录
训练	encryption_key	key	加密密钥
训练	model.pretrained_backbone_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
训练	results_dir	output_dir	当前作业结果目录
训练	train.pretrained_model_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
训练	train.resume_training_checkpoint_path	resume_model	从当前作业结果文件夹推断出的模型文件

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

部署

• tao-deploy-rtdetr