关于 tao-train-mask-grounding-dino
tao-train-mask-grounding-dino 是 NVIDIA TAO 中用于文本提示引导开放集实例分割的技能。它在 Grounding DINO 上加入掩码预测头,并支持训练、量化及 TensorRT 部署相关流程。
主要功能
- 01文本提示引导开放集分割
- 02Grounding DINO 掩码预测头
- 03训练、评估与推理
- 04模型量化与导出
- 05TensorRT 部署流程
适用场景
- 01根据文本提示分割图像对象
- 02训练带掩码标注的实例分割模型
- 03量化模型并运行 TensorRT 推理
安装
复制下面的命令,在终端中运行。
npx skills add https:/ / github.com/ promptingcompany/ nv-skills --skill tao-train-mask-grounding-dino安装后在 Claude Code 中输入/tao-train-mask-grounding-dino确认已加载。
包含的文件
除 SKILL.md 外,安装时一起下载的参考资料和脚本。
BENCHMARK.mdskill-card.mdskill.oms.sigevals/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-mask-grounding-dino.mdtao-deploy-mask-grounding-dino.skill_info.yamlschemas/evaluate.schema.jsonexport.schema.jsongen_trt_engine.schema.jsoninference.schema.jsonmanifest.jsonquantize.schema.jsontrain.schema.json详细介绍
Mask Grounding DINO
用于定位实例分割的 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。
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,但在生成架构之前无法运行 AutoML。
evaluate、inference、export 和部署流程等非训练操作仍由此模型技能处理。单次运行的 automl_policy 覆盖值不会更改模型元数据。
训练要求
- 数据集类型: 分割
- 格式: odvg、coco、coco_raw
- 监控指标: [bbox] val_mAP@50
各操作的数据集要求
| 操作 | Spec Key | 来源 | 文件 | 列表? |
|---|---|---|---|---|
| evaluate | dataset.test_data_sources | eval_dataset | image_dir: images.tar.gz、json_file: annotations.json | 否 |
| inference | dataset.infer_data_sources | inference_dataset | image_dir: images.tar.gz、classmap: label_map.txt、json_file: inference.jsonl、captions: inference.jsonl | 否 |
| 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 | 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 | 否 |
典型的规范覆盖项
每个操作都必须覆盖数据源——Agent 必须根据上方的“每个操作的数据集要求”表构建数据源路径,并将其包含在 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,
"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(必需数据源):
{
"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(必需数据源):
{
"infer_data_sources.data_type": "OD",
"dataset.infer_data_sources": {"image_dir": f"{S3_EVAL}/images.tar.gz", "classmap": f"{S3_EVAL}/label_map.txt", "json_file": f"{S3_EVAL}/inference.jsonl", "captions": f"{S3_EVAL}/inference.jsonl"},
}
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"},
}
评估数据集
可选。即使训练使用 ODVG,验证也会使用 COCO 格式的标注。
重要参数
- 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.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 管理。DDP/FSDP 行为与 Grounding DINO 相同。
| Spec Key | 说明 | 默认值 |
|---|---|---|
train.num_gpus | GPU 数量 | 1 |
train.gpu_ids | GPU 设备索引 | [0] |
train.num_nodes | 节点数量 | 1 |
train.distributed_strategy | ddp 或 fsdp | ddp |
硬件
至少需要 1 块 GPU,建议使用 4 块 GPU。每块 GPU 需配备 24GB 以上显存(推荐 A100)。由于增加了掩码预测头,资源需求高于 Grounding DINO。建议使用 24GB 以上显存的 GPU。
错误模式
CUDA 显存不足:降低 batch_size。掩码预测会在 Grounding DINO 的基础上增加额外开销。
Spec 参数/父模型推断
特定于模型的推断映射应放在此 MD 文件中,而不是 config.json 中。生成的运行脚本应读取本节内容,并在 create_job() 之前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程一致。
TAO Core mask_grounding_dino.config.json 中的推断映射:
| 操作 | Spec 字段 | 推断函数 | 含义 |
|---|---|---|---|
| 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_model 或 parent_model_folder,将上游 train/export/AutoML 子作业 ID 作为 parent_job_id 传入。SDK 会列出父结果文件夹中的内容,筛选检查点制品,并返回选定的模型文件或文件夹。不要将这些映射重新添加到 config.json 中,也不要修改生成的运行脚本来猜测检查点路径。
部署
- tao-deploy-mask-grounding-dino — 使用 TAO Deploy 生成 TensorRT 引擎、进行 TensorRT 评估和 TensorRT 推理的 Mask Grounding DINO 部署工作流。