关于 tao-train-grounding-dino
Grounding DINO 的 TAO 技能,面向需要训练和运行开放词汇目标检测模型的工作流。它结合视觉检测与 BERT 文本编码器,可按文本提示检测未预设类别的对象。
主要功能
- 01文本提示驱动的开放集检测
- 02TAO 模型训练与评估
- 03推理、导出与量化
- 04AutoML 训练策略支持
- 05TensorRT 部署流程参考
适用场景
- 01按自然语言描述检测对象
- 02在自定义数据集上训练检测器
- 03评估模型并导出 ONNX
- 04通过 TensorRT 运行推理
安装
复制下面的命令,在终端中运行。
npx skills add https:/ / github.com/ nvidia/ skills --skill tao-train-grounding-dino安装后在 Claude Code 中输入/tao-train-grounding-dino确认已加载。
包含的文件
除 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-grounding-dino.mdtao-deploy-grounding-dino.skill_info.yamlschemas/evaluate.schema.jsonexport.schema.jsongen_trt_engine.schema.jsoninference.schema.jsonmanifest.jsonquantize.schema.jsontrain.schema.json详细介绍
Grounding DINO
独立安装? 如果此会话不是由 TAO skill bank 插件初始化的,请先运行
tao-setup技能(主机预检、凭据、跨技能发现)。
Grounding DINO 用于开放词汇目标检测。它将 DINO 风格的检测与 BERT 文本编码器相结合,实现语言引导的检测。无需固定的类别词汇表,即可检测文本提示词所描述的对象。
完整 Grounding DINO 权重请设置 train.pretrained_model_path;仅使用骨干网络权重请设置 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/<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 已启用,但在为此模型生成架构之前无法运行。
evaluate、inference、export 和部署流程等非训练操作仍由此模型技能处理。每次运行的 automl_policy 覆盖值不会改变模型元数据。
训练要求
- 数据集类型: object_detection
- 格式: odvg、coco、raw
- 训练监控指标:
val_mAP50 - AutoML 指标约定: 对基于评估的选择,请使用最大化方向的
test_mAP50。打包的evaluate操作会输出test_mAP50; AutoML 试验比较和最终选择应使用此评估器 KPI, 而不是仅用于训练的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 | workflow prompts | 提示词列表 | 是 |
| 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 | calibration/eval dataset | 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 前解压归档文件。
常见规范覆盖项
每项操作都必须覆盖数据源——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"},
}
deploy/gen_trt_engine(使用 references/tao-deploy-grounding-dino.md):
{
"gen_trt_engine.onnx_file": "<exported_onnx_uri>",
"gen_trt_engine.trt_engine": "<output_engine_path>",
"gen_trt_engine.tensorrt.data_type": "FP16",
}
inference(必需的数据源):
{
"inference.checkpoint": "<selected train/AutoML checkpoint>",
"dataset.infer_data_sources.image_dir": [f"{S3_EVAL}/images.tar.gz"],
"dataset.infer_data_sources.captions": [
"fire extinguisher",
"cone",
"cart",
"forklift"
],
}
evaluate(必需的数据源):
{
"evaluate.checkpoint": "<selected train/AutoML checkpoint>",
"dataset.test_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}
quantize(必需的数据源):
{
"quantize.model_path": "<selected train checkpoint or exported ONNX model>",
"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"},
}
评估数据集
可选。尽管训练可以使用 ODVG 格式,验证仍使用 COCO 格式的标注来计算 mAP。
重要参数
- model.backbone:默认值为 swin_tiny_224_1k。还支持 resnet_50 和其他 Swin 变体。Swin 通常在定位任务中表现更好。
- model.text_encoder_type:用于文本编码的 BERT 模型。默认值为 bert-base-uncased。max_text_len 的默认值为 256。
- model.max_text_len:保持此值与数据集标签/token 位置映射一致。除非使用相同长度重新生成相应的标签映射,否则不要在冒烟测试中缩小此值;否则,token 概率与位置映射之间可能出现矩阵形状不匹配,导致验证失败。
- 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 目标数。在密集图像上,20 这类很小的冒烟测试值可能会导致匈牙利目标索引失败;除非确定数据集每张图像中的对象更少,否则 Grounding DINO 的最小冒烟测试应至少使用 100。 - train.optim.lr_steps:MultiStep 学习率调度。默认值为 [10]。
多 GPU/多节点
**启动方式:**由 Lightning 管理(单个 python 进程,由 Lightning 启动工作进程)。
| 规范键 | 说明 | 默认值 |
|---|---|---|
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 环境变量。
导出 / TRT 默认值
- 导出输入:960x544(比其他 OD 模型更大),opset 17。冒烟测试时,Grounding-DINO 导出规格应保持模板导出分辨率;将导出尺寸缩小到 128x128 等非常小的图像尺寸,可能会在
torch.onnx.export期间触发对比文本头中的 PyTorch ONNX 形状推断断言。 - 父级 PyTorch
grounding_dinoCLI 支持train、evaluate、inference、export和quantize。TensorRT 引擎生成、TensorRT 推理和 TensorRT 评估请参阅references/tao-deploy-grounding-dino.md。 - TRT 数据类型:仅支持 FP32、FP16 —— 不支持 INT8
- TRT 工作区:8192 MB(是其他 OD 模型的 8 倍)
- TRT max_batch_size:4
硬件
最低需要 1 块 GPU,建议使用 4 块 GPU。每块 GPU 需要 24GB 以上显存(推荐 A100)。由于 Grounding DINO 包含文本编码器(BERT),因此比标准 DINO 更重。建议使用 24GB 以上 GPU 显存。对于 16GB GPU,请减小 batch_size。
错误模式
CUDA 显存不足:减小 batch_size(4 -> 2 -> 1)。BERT 文本编码器会在视觉主干网络之外显著增加内存开销。
验证集标注类别 ID:验证集标注的类别 ID 应从 0 开始,以便正确计算损失。如有需要,请转换标注格式。
文本编码器加载错误:确保容器可以下载 bert-base-uncased 权重,或提供本地路径。
在 TAO Toolkit 7.0.0-rc-226 中使用 PyTorch 检查点进行量化失败:
容器中的 Grounding-DINO 量化脚本在加载检查点时传入
cap_lists=None,导致 post_process.py 出错。ONNX 量化使用导出的
ONNX 工件和 COCO 校准数据,但默认的 rc-226 PyTorch 镜像也缺少
modelopt.onnx.quantization 模块。请将其视为镜像 / SDK 阻碍,而非检查点解析器问题。
post_process.py 中出现 mat1 and mat2 shapes cannot be multiplied:文本
token 长度和标签位置映射不一致,通常是因为将 model.max_text_len 覆盖为低于默认值 256,而数据集标签映射仍使用长度为 256 的位置映射。请恢复 model.max_text_len,或使用相同长度重新生成标签映射。
criterion.py 中出现 index is out of bounds for dimension 0:对于当前批次中匹配的 ODVG 目标,model.num_queries
太小。请增大 model.num_queries,并确保 model.num_select 与之兼容。
使用 images.tar.gz/<image>.jpg 时出现 NotADirectoryError:直接调用 TAO CLI 时,程序试图将归档路径作为目录遍历。请解压归档,并将相关的 image_dir 字段设为解压后的图像文件夹;归档文件支持的技能数据源使用 runtime: extracted_folder,原因就在于此。
规格参数/父模型推断
特定模型的推断映射应放在此 MD 文件中,而不是 config.json 中。生成的运行器应读取本节,并在调用 create_job() 前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程相同。
TAO Core grounding_dino.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,也不要修改生成的运行器脚本来猜测检查点路径。
在 SDK 解析器之外选择 Grounding-DINO 检查点时,应精确匹配所需的 epoch/step 制品,例如
model_epoch_000_step_00046.pth。仅在明确请求 latest 时,gdino_model_latest.pth 符号链接才有效。将 model.backbone、model.num_queries、model.num_select、model.num_feature_levels、model.max_text_len 等结构性模型设置以及导出输入分辨率传递到 evaluate、inference、export 和 deploy 规格中,以确保检查点与引擎的形状匹配。