mcpField
T
tao-train-grounding-dino

tao-train-grounding-dino

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

关于 tao-train-grounding-dino

Grounding DINO 的 TAO 技能,面向需要训练和运行开放词汇目标检测模型的工作流。它结合视觉检测与 BERT 文本编码器,可按文本提示检测未预设类别的对象。

主要功能

  1. 01文本提示驱动的开放集检测
  2. 02TAO 模型训练与评估
  3. 03推理、导出与量化
  4. 04AutoML 训练策略支持
  5. 05TensorRT 部署流程参考

适用场景

  1. 01按自然语言描述检测对象
  2. 02在自定义数据集上训练检测器
  3. 03评估模型并导出 ONNX
  4. 04通过 TensorRT 运行推理

安装

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

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

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

包含的文件

除 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-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

独立安装? 如果此会话不是由 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 键。
各操作的数据集要求
操作规范键来源文件是否为列表?
evaluatedataset.test_data_sourceseval_datasetimage_dir: images.tar.gz, json_file: annotations.json否
inferencedataset.infer_data_sources.image_dirinference_datasetimages.tar.gz是
inferencedataset.infer_data_sources.captionsworkflow prompts提示词列表是
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_sourcescalibration/eval datasetimage_dir: images.tar.gz, json_file: annotations.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否

运行器可以将图像归档文件作为 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_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 环境变量。

导出 / TRT 默认值

  • 导出输入:960x544(比其他 OD 模型更大),opset 17。冒烟测试时,Grounding-DINO 导出规格应保持模板导出分辨率;将导出尺寸缩小到 128x128 等非常小的图像尺寸,可能会在 torch.onnx.export 期间触发对比文本头中的 PyTorch ONNX 形状推断断言。
  • 父级 PyTorch grounding_dino CLI 支持 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 中的推断映射:

操作规格字段推断函数含义
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,也不要修改生成的运行器脚本来猜测检查点路径。

在 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 规格中,以确保检查点与引擎的形状匹配。

部署

评论

支持 Markdown