mcpField
T
tao-train-grounding-dino

tao-train-grounding-dino

github.com/promptingcompany数据与数据库039 次安装17 周前更新
代码仓库

关于 tao-train-grounding-dino

用于训练和运行 Grounding DINO 开放集目标检测的 TAO skill。它通过文本提示进行语言引导检测,不依赖固定类别词表,并涵盖评估、推理、导出和部署流程。

主要功能

  1. 01文本提示开放集检测
  2. 02支持 ODVG、COCO 和 raw 格式
  3. 03训练、评估与推理流程
  4. 04ONNX 导出和 TensorRT 部署
  5. 05支持 AutoML 的训练策略

适用场景

  1. 01检测文本描述的对象
  2. 02使用 ODVG 或 COCO 标注训练
  3. 03在验证数据上评估检测结果
  4. 04通过 TensorRT 部署模型

安装

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

$npx skills add https://github.com/promptingcompany/nv-skills --skill tao-train-grounding-dino

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

包含的文件

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

BENCHMARK.md
skill-card.md
skill.oms.sig
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-grounding-dino.md
tao-deploy-grounding-dino.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

详细介绍

Grounding DINO

Grounding DINO 用于开放集目标检测。它将 DINO 风格的检测与 BERT 文本编码器相结合,实现语言引导的检测。无需固定类别词汇表,即可检测文本提示词描述的对象。

设置 train.pretrained_model_path 以使用完整的 Grounding DINO 权重,或设置 model.pretrained_backbone_path 以仅使用骨干网络权重。

对于 TAO Deploy TensorRT 操作(gen_trt_engine、TensorRT evaluate 和 TensorRT inference),请先阅读 references/tao-deploy-grounding-dino.md。部署规范模板位于此技能的 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/train.schema.json 和 references/spec_template_train.yaml。对于 automl_default_parameters、automl_disabled_parameters、默认值、最小值/最大值边界、枚举、选项权重、数学条件、依赖项和常用参数,请使用打包的训练架构。运行时不要依赖 ~/tao-core;维护者会在打包技能库之前重新生成架构/模板。

训练操作策略

此模型在模型层启用了 AutoML。处理任何训练阶段请求之前,请先阅读 references/skill_info.yaml,并根据明确指定的 automl_policy 值或用户的工作流请求确定本次运行的覆盖设置。将“关闭 AutoML”“禁用 AutoML”“不使用 HPO”或“普通训练”等说法视为仅对本次运行设置 automl_policy: off;否则默认使用 auto。当 automl_policy: auto、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 覆盖设置不会更改模型元数据。

训练要求

  • 数据集类型: object_detection
  • 格式: odvg、coco、raw
  • 监控指标: val_mAP50
各操作的数据集要求
操作规范键来源文件是否为列表?
evaluatedataset.test_data_sourceseval_datasetimage_dir: images.tar.gz, json_file: annotations.json否
inferencedataset.infer_data_sourcesinference_datasetimage_dir: images.tar.gz, classmap: label_map.txt否
quantizedataset.train_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json是
quantizedataset.val_data_sourceseval_datasetimage_dir: images.tar.gz, json_file: annotations.json否
quantizedataset.quant_calibration_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json否
traindataset.train_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json是
traindataset.val_data_sourceseval_datasetimage_dir: images.tar.gz, json_file: annotations.json否
常见 Spec 覆盖项

每项操作都必须覆盖数据源——Agent 必须根据上方的“每项操作的数据集要求”表构造数据源路径,并将其包含在 spec_overrides 中。

S3_TRAIN = "s3://bucket/data/train"
S3_EVAL = "s3://bucket/data/eval"

train(必需数据源):

{
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "train.num_gpus": 1,
    "dataset.train_data_sources": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations_odvg.jsonl", "label_map": f"{S3_TRAIN}/annotations_odvg_labelmap.json"}],
    "dataset.val_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

gen_trt_engine:

{
    "gen_trt_engine.tensorrt.data_type": "FP16",
}

inference(必需数据源):

{
    "dataset.infer_data_sources.captions": [
        "person"
    ],
    "dataset.infer_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "classmap": f"{S3_EVAL}/label_map.txt"},
}

evaluate(必需数据源):

{
    "dataset.test_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

quantize(必需数据源):

{
    "dataset.train_data_sources": [{"image_dir": f"{S3_TRAIN}/images.tar.gz", "json_file": f"{S3_TRAIN}/annotations_odvg.jsonl", "label_map": f"{S3_TRAIN}/annotations_odvg_labelmap.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_odvg.jsonl", "label_map": f"{S3_TRAIN}/annotations_odvg_labelmap.json"},
}

评估数据集

可选。验证使用 COCO 格式标注计算 mAP,即使训练可以使用 ODVG 格式。

重要参数

  • model.backbone:默认值为 swin_tiny_224_1k。也支持 resnet_50 和其他 Swin 变体。Swin 通常在 grounding 任务中的表现更好。
  • model.text_encoder_type:用于文本编码的 BERT 模型。默认值为 bert-base-uncased。max_text_len 的默认值为 256。
  • train.optim.lr:学习率。默认值为 2e-4。lr_backbone 为 2e-5。除 fp16/fp32 外,还支持 bf16 精度。
  • dataset.max_labels:训练期间每张图像的最大标签数。默认值为 50。对于密集标注数据集,可增大此值。
  • model.num_queries:对象查询数。由于采用开放词汇特性,默认值为 900(高于 DINO 的 300)。
  • train.optim.lr_steps:MultiStep 学习率调度。默认值为 [10]。

多 GPU / 多节点

**启动方式:**由 Lightning 管理(单个 python 进程,Lightning 会启动工作进程)。

Spec 键说明默认值
train.num_gpusGPU 数量1
train.gpu_idsGPU 设备索引[0]
train.num_nodes节点数量1
train.distributed_strategyddp 或 fsdpddp

与 DINO 的 DDP/FSDP 行为相同。多节点运行需要由编排器设置 WORLD_SIZE、NODE_RANK、MASTER_ADDR、MASTER_PORT 环境变量。

导出 / TRT 默认值

  • 导出输入:960x544(比其他 OD 模型更大),opset 17
  • TRT 数据类型:仅 FP32、FP16——不支持 INT8
  • TRT 工作区:8192 MB(比其他 OD 模型大 8 倍)
  • TRT max_batch_size:4

完整的 TAO Deploy 参考文档:tao-deploy-grounding-dino。

硬件

至少需要 1 块 GPU,建议使用 4 块 GPU。每块 GPU 需要 24GB 以上显存(推荐 A100)。由于包含文本编码器(BERT),Grounding DINO 比标准 DINO 更耗资源。建议使用显存 24GB 以上的 GPU。对于 16GB 显存的 GPU,请减小 batch_size。

错误模式

CUDA 显存不足:减小 batch_size(4 -> 2 -> 1)。除视觉主干网络外,BERT 文本编码器还会显著增加显存开销。

验证标注类别 ID:验证标注的类别 ID 应从 0 开始,以确保正确计算损失。如有需要,请使用标注格式转换。

文本编码器加载错误:确保容器可以下载 bert-base-uncased 权重,或提供本地路径。

规范参数 / 父模型推断

特定于模型的推断映射应放在此 MD 文件中,而不是放在 config.json 中。生成的运行器应读取本节内容,并在 create_job() 之前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程相同。

TAO Core grounding_dino.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当前作业的结果目录
gen_trt_engineencryption_keykey加密密钥
gen_trt_enginegen_trt_engine.onnx_fileparent_model从父作业结果文件夹推断的模型文件
gen_trt_enginegen_trt_engine.trt_enginecreate_engine_file输出 TensorRT 引擎路径
gen_trt_engineresults_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