TAOGroundingDINO训练技能Skill tao-train-grounding-dino

本技能提供 NVIDIA TAO 框架下 Grounding DINO 模型的训练、评估、导出、量化与部署能力。Grounding DINO 是一种开放集目标检测模型,结合 DINO 检测架构与 BERT 文本编码器,支持基于文本提示的开放词汇目标检测,无需固定类别。技能涵盖 AutoML 自动化训练、ODVG/COCO 数据格式、多 GPU/多节点训练、TensorRT FP16/FP32 推理等。关键词:Grounding DINO、开放词汇检测、语言引导检测、TAO、目标检测、AutoML、TensorRT、文本提示。

视觉模型训练 0 次安装 0 次浏览 更新于 9/6/2026
名称 tao-train-grounding-dino
描述 Grounding DINO 用于开放集目标检测。结合 DINO 风格检测与 BERT 文本编码器,实现语言引导检测——无需固定类别词汇,即可检测文本提示所描述的对象。在针对 TAO Grounding DINO 模型进行训练、评估、导出、量化或运行推理时使用。触发短语包括“train Grounding DINO”、“open-vocabulary detection”、“text-prompted detector”、“language-guided object detection”。
开源协议 Apache-2.0 compatibility: 需要 docker + nvidia-container-toolkit。 metadata:
版本 0.1.0
作者 NVIDIA Corporation allowed-tools: Read Bash tags: - object - detection

Grounding DINO

需要独立安装?如果此会话未由 TAO 技能库插件初始化,请先运行 tao-setup 技能(主机预检、凭据、跨技能发现)。

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。

数据类模式

生成的 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 覆盖不会改变模型元数据。

训练要求

  • 数据集类型: object_detection
  • 格式: odvg、coco、raw
  • 监控指标: val_mAP50

逐操作的数据集要求

操作 规范键 来源 文件 列表?
evaluate dataset.test_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json
inference dataset.infer_data_sources.image_dir inference_dataset images.tar.gz
inference dataset.infer_data_sources.captions 工作流提示 提示列表
quantize dataset.train_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json
quantize dataset.val_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json
quantize dataset.quant_calibration_data_sources 校准/评估数据集 image_dir: images.tar.gz, json_file: annotations.json
train dataset.train_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json
train dataset.val_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json

运行器可以将图像存档作为 images.tar.gz 提供,但直接的本地 Docker TAO CLI 规范必须将 image_dir 指向提取的图像目录。技能元数据将这些存档支持的图像源标记为 runtime: extracted_folder,以便新运行器在启动 TAO 前解包存档。

典型的规范覆盖

数据源覆盖对于每个操作都是强制性的——代理必须按照上面的逐操作数据集要求表构造数据源路径,并将它们包含在 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,
    '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(使用 references/tao-deploy-grounding-dino.md):

{
    'gen_trt_engine.onnx_file': '<已导出的 onnx uri>',
    'gen_trt_engine.trt_engine': '<输出引擎路径>',
    'gen_trt_engine.tensorrt.data_type': 'FP16',
}

推理(强制数据源):

{
    'inference.checkpoint': '<选定的训练/AutoML 检查点>',
    'dataset.infer_data_sources.image_dir': [f'{S3_EVAL}/images.tar.gz'],
    'dataset.infer_data_sources.captions': [
        'fire extinguisher',
        'cone',
        'cart',
        'forklift'
    ],
}

评估(强制数据源):

{
    'evaluate.checkpoint': '<选定的训练/AutoML 检查点>',
    'dataset.test_data_sources': {'image_dir': f'{S3_EVAL}/images.tar.gz', 'json_file': f'{S3_EVAL}/annotations.json'},
}

量化(强制数据源):

{
    'quantize.model_path': '<选定的训练检查点或导出的 ONNX 模型>',
    '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_EVAL}/images.tar.gz', 'json_file': f'{S3_EVAL}/annotations.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。
  • model.max_text_len:保持与数据集标签/标记位置映射对齐。除非对应的标签映射以相同长度重新生成,否则不要为了冒烟测试缩小它。
  • train.optim.lr:学习率。默认 2e-4。lr_backbone 2e-5。除 fp16/fp32 外还支持 bf16 精度。
  • dataset.max_labels:训练期间每张图像的最大标签数。默认 50。密集注释数据集请增大。
  • model.num_queries:对象查询数。默认 900,因开放词汇特性高于 DINO 的 300。
  • model.num_queries / model.num_select:保持 num_queries 足够高以匹配批中匹配的 ODVG 目标数。最小冒烟运行建议至少 100。
  • train.optim.lr_steps:MultiStep 学习率调度。默认 [10]。

多 GPU / 多节点

启动方式:Lightning 托管(单个 python 进程,Lightning 生成 worker)。

规范键 描述 默认
train.num_gpus GPU 数量 1
train.gpu_ids GPU 设备索引 [0]
train.num_nodes 节点数量 1
train.distributed_strategy ddp 或 fsdp ddp

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

导出 / TRT 默认值

  • 导出输入:960x544,比其他 OD 模型大,opset 17。请保持 Grounding-DINO 导出规范在模板导出分辨率下进行冒烟测试。
  • 父级 PyTorch grounding_dino CLI 支持 train、evaluate、inference、export 和 quantize。通过 references/tao-deploy-grounding-dino.md 运行 TensorRT 引擎生成、推理和评估。
  • TRT 数据类型:仅 FP32、FP16,不支持 INT8。
  • TRT 工作空间:8192 MB(比其他 OD 模型大 8 倍)。
  • TRT 最大批大小:4。

硬件

最低 1 GPU,推荐 4 GPU。每 GPU 需要24GB+ VRAM,建议 A100。Grounding DINO 因 BERT 文本编码器比标准 DINO 更重。对于 16GB GPU,请减小批大小。

错误模式

  • CUDA 内存不足:减小批大小 4 -> 2 -> 1。BERT 编码器增加大量内存开销。
  • 验证注释类别 ID 应从 0 开始。
  • 文本编码器加载错误:确保容器能访问 bert-base-uncased 权重或提供本地路径。
  • TAO Toolkit 7.0.0-rc-226 量化检查点失败:镜像缺少 modelopt.onnx.quantization 模块。
  • post_process.py 形状不匹配:恢复 model.max_text_len 或重新生成 label map。
  • criterion.py 索引越界:增大 model.num_queries。
  • NotADirectoryError:解压 images.tar.gz 并将 image_dir 指向目录。

规范参数 / 父模型推断

模型特定推断映射属于此 MD 文件,不使用 config.json。请使用 SDK helpers 应用映射。详细映射略。

选择检查点时,精确匹配 epoch/step 工件,如 model_epoch_000_step_00046.pth。传递结构化模型设置以保证形状一致。

部署

  • tao-deploy-grounding-dino (references/tao-deploy-grounding-dino.md)