掩码接地DINO(MaskGroundingDINO)Skill tao-train-mask-grounding-dino

该技能面向TAO框架中的Mask Grounding DINO模型,用于文本引导的开放词汇实例分割。它封装了模型训练、评估、推理、量化、导出和TensorRT部署等工作流,支持AutoML、多GPU/多节点分布式训练,并适配ODVG/COCO等分割数据集。关键词:Mask Grounding DINO、开放词汇分割、接地实例分割、文本提示分割、掩码预测、模型训练、模型评估、模型推理、模型量化、模型导出、TensorRT、TAO、AutoML、多GPU训练、分割数据集。

视觉模型训练 0 次安装 0 次浏览 更新于 9/6/2026
名称 tao-train-mask-grounding-dino
描述 Mask Grounding DINO 用于接地实例分割(grounded instance segmentation)。通过增加掩码预测头扩展 Grounding DINO,实现文本提示引导的开放集分割。训练、评估、导出、量化或运行 TAO Mask-Grounding-DINO 模型推理时使用。触发短语包括“train Mask Grounding DINO”、“open-vocabulary segmentation”、“text-prompted instance segmentation”、“grounded mask DETR”。
开源协议 Apache-2.0 compatibility: 需要 docker + nvidia-container-toolkit。 metadata:
版本 “0.1.0”
作者 NVIDIA Corporation allowed-tools: 读取 Bash tags: - segmentation

Mask Grounding DINO

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

Mask Grounding DINO 用于接地实例分割。通过掩码预测头扩展 Grounding DINO,支持文本提示引导的开放集分割。

设置 train.pretrained_model_path 以指定完整模型权重。

对于 TAO Deploy TensorRT 操作(gen_trt_engine、TensorRT evaluate、TensorRT inference),请先阅读 references/tao-deploy-mask-grounding-dino.md。部署规范模板位于本技能的 references/ 文件夹中,文件名前缀为 spec_template_deploy_*.yaml

数据类 Schema

生成的 TAO Core Schema 位于 schemas/<action>.schema.json,并且 schemas/manifest.json 列出了可用操作。每个生成的 Schema 也通过 Schema 顶层的 default 字段输出 references/spec_template_<action>.yaml。AutoML 的启用与否在 references/skill_info.yaml 的模型层通过 automl_enabled 声明。对于某个操作的可运行 AutoML,需要 schemas/<action>.schema.jsonreferences/spec_template_<action>.yaml 存在且可解析。对 automl_default_parametersautoml_disabled_parameters、默认值、最小值/最大值边界、枚举、选项权重、数学条件、依赖关系和常用参数,请使用打包的所选操作 schema。运行时不要依赖 ~/tao-core;维护人员会在打包技能库之前重新生成 schema/模板。

训练操作策略

本模型在模型层启用 AutoML。在处理任何训练阶段请求之前,请读取 references/skill_info.yaml,并从显式的 automl_policy 值或用户的工作流请求中解析运行覆盖。默认使用 automl_policy: on,并在新的启动提示中只暴露 on / off。将诸如“turn off AutoML”、“disable AutoML”、“no HPO”、“plain training”等说法视为本次运行的 automl_policy: off。当 automl_policy: onautoml_enabled: true 并且已打包了 schemas/train.schema.jsonreferences/spec_template_train.yaml 时,默认通过 tao-skill-bank:tao-run-automl 使用此模型的 skill_dir 路由训练操作。保留工作流/应用程序对数据集、规范、输出目录、GPU/平台设置、父检查点和 automl_policy 的覆盖。仅当 automl_policy: off 或打包的训练 schema/模板缺失时,才直接使用模型训练;在缺失 schema 的情况下,报告 AutoML 已启用但无法运行,直到生成 schema。

评估、推理、导出和部署流程等非训练操作保留在本模型技能中。每次运行的 automl_policy 覆盖不会改变模型元数据。

训练要求

  • 数据集类型: 分割(segmentation)
  • 格式: odvg, coco, coco_raw
  • 监控指标: val_loss

每个操作的数据集要求

操作 规格键 来源 文件 是否列表?
evaluate dataset.test_data_sources eval_dataset image_dir: images.tar.gz, json_file: annotations.json
evaluate dataset.test_data_sources.data_type eval_dataset OD
inference dataset.infer_data_sources inference_dataset image_dir: images.tar.gz, captions: 文本提示
inference dataset.infer_data_sources.data_type inference_dataset OD
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.val_data_sources.data_type eval_dataset OD
quantize dataset.quant_calibration_data_sources train_datasets image_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.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
train dataset.val_data_sources.data_type eval_dataset OD

典型规格覆盖

数据源覆盖对于每个操作都是强制性的——代理必须根据上面的每个操作数据集要求表构建数据源路径,并将它们包含在 spec_overrides 中。

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

train(强制数据源):

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "dataset.val_data_sources.data_type": "OD",
    "model.num_region_queries": 100,
    "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"},
}

evaluate(强制数据源):

{
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.test_data_sources.data_type": "OD",
    "dataset.test_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "json_file": f"{S3_EVAL}/annotations.json"},
}

inference(强制数据源):

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.infer_data_sources.data_type": "OD",
    "dataset.infer_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "captions": ["person", "bicycle", "car"]},
}

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 格式的注释,即使训练使用 ODVG。

重要参数

  • model.backbone:默认 swin_tiny_224_1k。与 Grounding DINO 相同的骨干网络选项。
  • train.optim.lr:学习率。默认 2e-4。lr_backbone 2e-5。复用 GDINOTrainExpConfig——与 Grounding DINO 相同的训练设置。
  • model.num_queries:目标查询数量。默认 900。
  • model.enc_layers / model.dec_layers:训练/AutoML 运行中都保持为 6。掩码头在验证期间断言六个解码器输出,因此复制 Grounding DINO smoke 覆盖以减少 transformer 层会导致立即失败。
  • AutoML 指标备注:对于训练阶段 AutoML,使用 metric="val_loss"direction="minimize"。打包的训练循环记录验证损失标量;训练作业期间不会发出 [bbox] val_mAP@50
  • model.has_mask:启用掩码预测头。默认为 True。增加 mask/dice/rela 损失系数。
  • model.num_region_queries:用于掩码预测的区域查询数量。默认 100。
  • model.loss_types:损失组件。默认为 [labels, boxes, masks]。包括 mask_loss_coef, dice_loss_coef, rela_loss_coef。
  • evaluate.ioi_threshold:掩码评估的 IoI 阈值。默认 0.5。
  • evaluate.nms_threshold:NMS 阈值。默认 0.2。
  • evaluate.text_threshold:文本匹配阈值。默认 0.3。
  • dataset.has_mask:数据集包含掩码注释。默认 True。val_data_sources 默认 data_type 为 “VG”。

多 GPU / 多节点

启动方式: Lightning 管理。与 Grounding DINO 相同的 DDP/FSDP 行为。

规格键 说明 默认值
train.num_gpus GPU 数量 1
train.gpu_ids GPU 设备索引 [0]
train.num_nodes 节点数量 1
train.distributed_strategy ddpfsdp ddp

硬件

最少 1 个 GPU,推荐 4 个 GPU。每个 GPU 24GB+(推荐 A100)显存。由于掩码预测头,比 Grounding DINO 更重。推荐 24GB+ GPU 内存。

错误模式

CUDA 内存不足:减小批次大小。掩码预测使模型在 Grounding DINO 之上增加了开销。

部署 schema 错误 test_threshold:TAO Deploy 使用 evaluate.text_thresholdinference.text_threshold。不要在部署规范中使用 test_threshold

部署模型形状不匹配:将 transformer 和掩码结构字段从导出阶段带入部署评估/推理规范,包括 model.num_queriesmodel.num_selectmodel.max_text_lenmodel.num_region_queriesmodel.has_mask。这些值必须与用于构建 TensorRT 引擎的 ONNX 模型匹配。

规格参数 / 父模型推断

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

来自 TAO Core mask_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 当前作业结果目录
gen_trt_engine encryption_key key 加密密钥
gen_trt_engine gen_trt_engine.onnx_file parent_model 从父作业结果文件夹推断出的模型文件
gen_trt_engine gen_trt_engine.trt_engine create_engine_file 输出 TensorRT 引擎路径
gen_trt_engine 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_modelparent_model_folder,请将上游 train/export/AutoML 子作业 ID 作为 parent_job_id 传递。SDK 列出父结果文件夹,筛选检查点工件,并返回所选模型文件或文件夹。不要将这些映射添加回 config.json,也不要修补生成的运行程序脚本来猜测检查点路径。

在 SDK 解析器之外选择 Mask Grounding DINO 检查点时,请精确匹配目标周期/步骤工件,例如 model_epoch_000_step_00049.pthmask_gdino_model_latest.pth 符号链接仅在明确要求最新版本时有效。父 PyTorch mask_grounding_dino CLI 支持 trainevaluateinferenceexportquantize;通过 references/tao-deploy-mask-grounding-dino.md 运行 TensorRT 引擎生成、TensorRT 推理和 TensorRT 评估。

部署