mcpField
T
tao-train-mask-grounding-dino

tao-train-mask-grounding-dino

github.com/promptingcompany数据与数据库039 次安装17 周前更新
代码仓库

关于 tao-train-mask-grounding-dino

tao-train-mask-grounding-dino 是 NVIDIA TAO 中用于文本提示引导开放集实例分割的技能。它在 Grounding DINO 上加入掩码预测头,并支持训练、量化及 TensorRT 部署相关流程。

主要功能

  1. 01文本提示引导开放集分割
  2. 02Grounding DINO 掩码预测头
  3. 03训练、评估与推理
  4. 04模型量化与导出
  5. 05TensorRT 部署流程

适用场景

  1. 01根据文本提示分割图像对象
  2. 02训练带掩码标注的实例分割模型
  3. 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.md
skill-card.md
skill.oms.sig
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-mask-grounding-dino.md
tao-deploy-mask-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

详细介绍

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来源文件列表?
evaluatedataset.test_data_sourceseval_datasetimage_dir: images.tar.gz、json_file: annotations.json否
inferencedataset.infer_data_sourcesinference_datasetimage_dir: images.tar.gz、classmap: label_map.txt、json_file: inference.jsonl、captions: inference.jsonl否
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_sourcestrain_datasetsimage_dir: images.tar.gz、json_file: annotations_odvg.jsonl、label_map: annotations_odvg_labelmap.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否
典型的规范覆盖项

每个操作都必须覆盖数据源——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_gpusGPU 数量1
train.gpu_idsGPU 设备索引[0]
train.num_nodes节点数量1
train.distributed_strategyddp 或 fsdpddp

硬件

至少需要 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 字段推断函数含义
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当前作业的结果目录
gen_trt_engineencryption_keykey加密密钥
gen_trt_enginegen_trt_engine.onnx_fileparent_model从父作业结果文件夹推断出的模型文件
gen_trt_enginegen_trt_engine.trt_enginecreate_engine_file输出 TensorRT 引擎路径
gen_trt_engineresults_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 中,也不要修改生成的运行脚本来猜测检查点路径。

部署

  • tao-deploy-mask-grounding-dino — 使用 TAO Deploy 生成 TensorRT 引擎、进行 TensorRT 评估和 TensorRT 推理的 Mask Grounding DINO 部署工作流。

评论

支持 Markdown