mcpField
T
tao-train-deformable-detr

tao-train-deformable-detr

github.com/nvidia开发3.8k1.6k 次安装今天更新
代码仓库

关于 tao-train-deformable-detr

此 TAO 技能用于训练和运行 Deformable DETR,面向二维目标检测,并通过可变形注意力处理多尺度特征。它支持 AutoML、检查点评估、推理、导出、量化以及由 TAO Deploy 执行的 TensorRT 引擎生成。

主要功能

  1. 01二维目标检测
  2. 02可变形注意力处理多尺度特征
  3. 03AutoML 训练支持
  4. 04评估、推理和导出
  5. 05量化与 TensorRT 部署

适用场景

  1. 01在自定义检测数据集上训练
  2. 02评估检查点并运行检测推理
  3. 03导出模型并生成 TensorRT 引擎

安装

复制下面的命令,在终端中运行。

$npx skills add https://github.com/nvidia/skills --skill tao-train-deformable-detr

安装后在 Claude Code 中输入/tao-train-deformable-detr确认已加载。

包含的文件

除 SKILL.md 外,安装时一起下载的参考资料和脚本。

BENCHMARK.md
skill-card.md
skill.oms.sig
config/
skillspector-baseline.yaml
evals/
evals.json
references/
skill_info.yaml
spec_template_deploy_evaluate.yaml
spec_template_deploy_gen_trt_engine.yaml
spec_template_deploy_inference.yaml
spec_template_evaluate.yaml
spec_template_export.yaml
spec_template_gen_trt_engine.yaml
spec_template_inference.yaml
spec_template_quantize.yaml
spec_template_train.yaml
tao-deploy-deformable-detr.md
tao-deploy-deformable-detr.skill_info.yaml
schemas/
evaluate.schema.json
export.schema.json
gen_trt_engine.schema.json
inference.schema.json
manifest.json
quantize.schema.json
train.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_direval_datasetimages.tar.gz否
评估dataset.test_data_sources.json_fileeval_datasetannotations.json否
导出dataset.train_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations.json是
导出dataset.val_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations.json是
推理dataset.infer_data_sources.image_dirinference_datasetimages.tar.gz是
推理dataset.infer_data_sources.classmapinference_datasetlabel_map.txt否
量化dataset.train_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations.json是
量化dataset.val_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations.json是
量化dataset.quant_calibration_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations.json否
训练dataset.train_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations.json是
训练dataset.val_data_sourcestrain_datasetsimage_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_gpusGPU 数量1
train.gpu_idsGPU 设备索引[0]
train.num_nodes节点数量1
train.distributed_strategyddp 或 fsdpddp

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 的推断映射:

操作规范字段推断函数含义
evaluateencryption_keykey加密密钥
evaluateevaluate.checkpointparent_model从父作业结果文件夹推断出的模型文件
evaluateevaluate.trt_engineparent_model从父作业结果文件夹推断出的模型文件
evaluateresults_diroutput_dir当前作业结果目录
exportencryption_keykey加密密钥
exportexport.checkpointparent_model从父作业结果文件夹推断出的模型文件
exportexport.onnx_filecreate_onnx_file输出 ONNX 路径
exportresults_diroutput_dir当前作业结果目录
inferenceencryption_keykey加密密钥
inferenceinference.checkpointparent_model从父作业结果文件夹推断出的模型文件
inferenceinference.trt_engineparent_model从父作业结果文件夹推断出的模型文件
inferenceresults_diroutput_dir当前作业结果目录
quantizeencryption_keykey加密密钥
quantizequantize.model_pathparent_model从父作业结果文件夹推断出的模型文件
quantizeresults_diroutput_dir当前作业结果目录
trainencryption_keykey加密密钥
trainmodel.pretrained_backbone_pathptm_if_no_resume_model不存在恢复检查点时使用的 PTM
trainresults_diroutput_dir当前作业结果目录
traintrain.pretrained_model_pathptm_if_no_resume_model不存在恢复检查点时使用的完整模型 PTM
traintrain.resume_training_checkpoint_pathresume_model从当前作业结果文件夹推断出的模型文件

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

部署

评论

支持 Markdown