关于 tao-train-deformable-detr
此 TAO 技能用于训练和运行 Deformable DETR,面向二维目标检测,并通过可变形注意力处理多尺度特征。它支持 AutoML、检查点评估、推理、导出、量化以及由 TAO Deploy 执行的 TensorRT 引擎生成。
主要功能
- 01二维目标检测
- 02可变形注意力处理多尺度特征
- 03AutoML 训练支持
- 04评估、推理和导出
- 05量化与 TensorRT 部署
适用场景
- 01在自定义检测数据集上训练
- 02评估检查点并运行检测推理
- 03导出模型并生成 TensorRT 引擎
安装
复制下面的命令,在终端中运行。
npx skills add https:/ / github.com/ nvidia/ skills --skill tao-train-deformable-detr安装后在 Claude Code 中输入/tao-train-deformable-detr确认已加载。
包含的文件
除 SKILL.md 外,安装时一起下载的参考资料和脚本。
BENCHMARK.mdskill-card.mdskill.oms.sigconfig/skillspector-baseline.yamlevals/evals.jsonreferences/skill_info.yamlspec_template_deploy_evaluate.yamlspec_template_deploy_gen_trt_engine.yamlspec_template_deploy_inference.yamlspec_template_evaluate.yamlspec_template_export.yamlspec_template_gen_trt_engine.yamlspec_template_inference.yamlspec_template_quantize.yamlspec_template_train.yamltao-deploy-deformable-detr.mdtao-deploy-deformable-detr.skill_info.yamlschemas/evaluate.schema.jsonexport.schema.jsongen_trt_engine.schema.jsoninference.schema.jsonmanifest.jsonquantize.schema.jsontrain.schema.json详细介绍
Deformable DETR
独立安装? 如果此会话不是由 TAO skill bank 插件初始化的,请先运行
tao-setup技能(主机预检、凭据、跨技能发现)。
用于 2D 目标检测的 Deformable DETR。使用可变形注意力,高效处理多尺度特征。与 DINO 相比更轻量,同时准确率具有竞争力。
使用预训练权重。设置 model.pretrained_backbone_path 以仅加载主干网络,或设置 train.pretrained_model_path 以初始化完整模型。
支持的父模型操作包括 train、evaluate、inference、export 和 quantize。此网络的 PyT 模型容器不支持原生的 gen_trt_engine 子任务。references/skill_info.yaml 中声明的 gen_trt_engine 操作必须使用 TAO Deploy 容器运行。Deploy 规范模板位于此技能的 references/ 文件夹中,文件名前缀为 spec_template_deploy_*.yaml。
Dataclass 架构
生成的 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;维护者会在打包 skill bank 前重新生成架构和模板。
训练操作策略
此模型在模型层级启用了 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,但在生成架构之前无法运行 AutoML。
evaluate、inference、export 等非训练操作以及部署流程仍由此模型技能处理。每次运行的 automl_policy 覆盖设置不会更改模型元数据。
训练要求
- 数据集类型: object_detection
- 格式: coco、coco_raw
- AutoML 指标约定: 对于基于评估结果的选择,使用
test_mAP50,并设为最大化方向。仅在只依据训练日志的 AutoML 工作流中使用val_mAP50。 - 训练监控指标:
val_mAP50用于 AP50;val_mAP用于 COCO/论文风格的基准比较。 - Evaluate 操作指标:
test_mAP50用于 AP50;test_mAP用于 COCO/论文风格的基准比较。通过 evaluate 操作比较的 AutoML 试验必须使用相应的test_*KPI,而不是在评估器输出中查找val_*KPI。
各操作的数据集要求
| 操作 | 规格键 | 来源 | 文件 | 列表? |
|---|---|---|---|---|
| 评估 | dataset.test_data_sources.image_dir | eval_dataset | images.tar.gz | 否 |
| 评估 | dataset.test_data_sources.json_file | eval_dataset | annotations.json | 否 |
| 导出 | dataset.train_data_sources | train_datasets | image_dir: images.tar.gz, json_file: annotations.json | 是 |
| 导出 | dataset.val_data_sources | train_datasets | image_dir: images.tar.gz, json_file: annotations.json | 是 |
| 推理 | dataset.infer_data_sources.image_dir | inference_dataset | images.tar.gz | 是 |
| 推理 | dataset.infer_data_sources.classmap | inference_dataset | label_map.txt | 否 |
| 量化 | dataset.train_data_sources | train_datasets | image_dir: images.tar.gz, json_file: annotations.json | 是 |
| 量化 | dataset.val_data_sources | train_datasets | 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 | train_datasets | image_dir: images.tar.gz, json_file: annotations.json | 是 |
常见规格覆盖项
每个操作都必须覆盖数据源——Agent 必须根据上方的各操作数据集要求表构建数据源路径,并将其包含在 spec_overrides 中。
S3_TRAIN = "s3://bucket/data/train"
S3_EVAL = "s3://bucket/data/eval"
训练(必须提供数据源):
{
"train.num_epochs": 10,
"train.checkpoint_interval": 10,
"train.validation_interval": 10,
"train.num_gpus": 1,
"train.gpu_ids": [0],
"dataset.num_classes": "<object classes> + 1",
"dataset.eval_class_ids": [1, 2, "..."],
"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_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
}
评估(必须提供数据源):
{
"dataset.num_classes": "<object classes> + 1",
"dataset.eval_class_ids": [1, 2, "..."],
"dataset.test_data_sources.image_dir": f"{S3_EVAL}/images.tar.gz",
"dataset.test_data_sources.json_file": f"{S3_EVAL}/annotations.json",
}
如果训练或 AutoML 运行更改了影响架构的字段,例如
model.enc_layers、model.dec_layers、model.num_queries 或
model.num_select,则在使用所选检查点执行评估、导出、推理和
部署操作时,也要传入相同的值。除了上述字段外,如果
model.num_feature_levels、model.dim_feedforward、输入图像
尺寸和数据集类别元数据曾被更改,也要一并传入。将检查点加载到默认架构
可能会因张量形状不匹配而失败,在为提高速度而缩小 Transformer 规模的
冒烟测试运行中尤其如此。
导出(必须提供数据源):
{
"dataset.num_classes": "<object classes> + 1",
"dataset.eval_class_ids": [1, 2, "..."],
"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_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
}
TensorRT 引擎生成:
export 后使用 deploy 规格模板。不要从父级 PyT 模型容器调用
deformable_detr gen_trt_engine;该 CLI 声明支持
convert、evaluate、export、inference、quantize、train 和
default_specs,但不支持 gen_trt_engine。模型操作元数据会选择
TAO Deploy 容器来生成引擎。
Deploy 引擎生成需要将导出的 ONNX 文件作为输入,并在
gen_trt_engine.trt_engine 处创建引擎。
{
"gen_trt_engine.tensorrt.data_type": "FP16",
"dataset.num_classes": "<object classes> + 1",
"gen_trt_engine.tensorrt.calibration.cal_image_dir": [f"{S3_TRAIN}/images.tar.gz"],
}
推理(必须提供数据源):
{
"dataset.num_classes": "<object classes> + 1",
"dataset.infer_data_sources.image_dir": [f"{S3_EVAL}/images.tar.gz"],
"dataset.infer_data_sources.classmap": f"{S3_EVAL}/label_map.txt",
}
量化(必须提供数据源):
{
"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_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"}],
"dataset.quant_calibration_data_sources": {"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations.json"},
}
评估数据集
可选。如果提供了数据集,则会在每个检查点间隔计算验证集 mAP。
检查点处理
训练会按 model_epoch_<epoch>_step_<step>.pth 模式生成带有 epoch 和 step 信息的检查点,
并创建 dd_model_latest.pth 符号链接。对于有依赖关系的操作,请使用模型专用或 SDK 提供的检查点解析器,
选择目标产物。除非用户明确要求使用最新检查点,否则评估、推理、导出和量化应接收所选检查点的确切路径,而不是 dd_model_latest.pth 符号链接。
恢复或重新训练时,应将 train.resume_training_checkpoint_path 设置为要恢复训练的确切检查点路径。
重要参数
- dataset.num_classes:目标类别数加上背景类别。默认值为 91(COCO)。必须与标注一致。
- dataset.eval_class_ids:纳入 COCO 指标的前景类别 ID。对于自定义数据集,请将其设置为所有目标类别 ID;模板默认仅评估类别 ID 1。
- model.backbone:默认值为 resnet_50。支持:resnet_50、gcvit_tiny、gcvit_small、gcvit_base、gcvit_large、gcvit_large_384(支持范围比 DINO 更有限)。
- train.optim.lr:学习率。默认值为 2e-4(AdamW)。lr_backbone 为 2e-5。
- train.optim.lr_steps:MultiStep 学习率调度。默认值为 [40]。对于短周期训练,请将其设置为总 epoch 数的约 80%。
- model.num_queries:目标查询数量。默认值为 300。有效范围为 100-900。
- model.dropout_ratio:Transformer 层中的 Dropout。默认值为 0.3(高于 DINO 的 0.0)。数据集较大时调低,较小时调高。
- model.dim_feedforward:FFN 隐藏维度。默认值为 1024(DINO 为 2048)。增大此值可以提升容量,但会增加内存开销。
多 GPU / 多节点
**启动方式:**由 Lightning 管理(单个 python 进程,由 Lightning 启动工作进程)。
| Spec Key | 说明 | 默认值 |
|---|---|---|
train.num_gpus | GPU 数量 | 1 |
train.gpu_ids | GPU 设备索引 | [0] |
train.num_nodes | 节点数量 | 1 |
train.distributed_strategy | ddp 或 fsdp | ddp |
DDP/FSDP 行为与 DINO 相同。多节点运行需要由编排器设置 WORLD_SIZE、NODE_RANK、MASTER_ADDR、MASTER_PORT 环境变量。
增大 train.num_gpus 时,还要将 train.gpu_ids 设置为相同的可见
设备范围。例如,单节点 8 GPU 的 Slurm 运行必须同时包含
"train.num_gpus": 8 和 "train.gpu_ids": [0, 1, 2, 3, 4, 5, 6, 7]。
导出 / TRT 默认值
- 导出输入:640x640,opset 17
- TRT 数据类型:FP32、FP16、INT8
- TRT 工作区:1024 MB
- TRT max_batch_size:1
硬件
至少需要 1 个 GPU,建议使用 4 个 GPU。每个 GPU 需要 16GB 以上显存(V100 或 A100)。由于 FFN 较小,所需资源略低于 DINO。batch_size=4 可在大多数 16GB 以上显存的 GPU 上运行。
错误模式
CUDA 显存不足:减小 batch_size(4 -> 2 -> 1)。
num_select must be < num_queries * num_classes:与 DINO 的限制相同。
return_interm_indices length must match num_feature_levels:默认值为 [1,2,3,4],num_feature_levels=4。
数据集大小小于总 batch 大小:减小 batch_size 或 num_gpus。
AutoML 指标提取:Deformable DETR 会在结构化训练状态中输出 val_mAP 和 val_mAP50,而独立的 evaluate 操作会输出 test_mAP 和 test_mAP50。对于 COCO/论文风格的仅训练比较,请优化 val_mAP;对于明确使用 AP50 的仅训练工作流,请优化 val_mAP50。当 AutoML 使用 evaluate 操作为每条建议评分时,请使用对应的 test_mAP 或 test_mAP50 KPI。优先使用 results_dir/train/status.json 或 AutoML 结果状态,再解析原始日志。默认检测模型调用时不要优化 val_loss。
规范参数 / 父模型推断
特定模型的推断映射应写在此 MD 文件中,而不是 config.json。生成的运行器应读取本节,并在调用 create_job() 之前使用 SDK 助手应用这些映射。这与旧版微服务的 infer_params.py 流程一致。
来自 TAO Core deformable_detr.config.json 的推断映射:
| 操作 | 规范字段 | 推断函数 | 含义 |
|---|---|---|---|
| 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,将上游 train/export/AutoML 子作业 ID 作为 parent_job_id 传入。SDK 会列出父作业结果文件夹,筛选检查点制品,并返回选定的模型文件或文件夹。不要将这些映射加回 config.json,也不要修改生成的运行器脚本来猜测检查点路径。