mcpField
T
tao-train-mask-grounding-dino

tao-train-mask-grounding-dino

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

关于 tao-train-mask-grounding-dino

此 TAO 技能用于训练和运行 Mask Grounding DINO,以文本提示引导开放类别实例分割。它扩展 Grounding DINO,加入掩码预测头,并涵盖训练、评估、推理及部署流程。

主要功能

  1. 01文本提示引导开放类别分割
  2. 02实例掩码预测
  3. 03训练支持 AutoML
  4. 04支持评估与推理
  5. 05支持导出和部署流程

适用场景

  1. 01按文本类别提示分割图像对象
  2. 02在分割数据集上训练模型
  3. 03评估检查点并进行推理

安装

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

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

安装后在 Claude Code 中输入/tao-train-mask-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-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

单独安装? 如果此会话并非由 TAO skill bank 插件初始化,请先运行 tao-setup skill(主机预检、凭据、跨 skill 发现)。

用于开放词汇实例分割的 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。Deploy 规格模板位于此 skill 的 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 前重新生成架构和模板。

Train 操作策略

此模型在模型层启用了 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 执行 train 操作,并使用此模型的 skill_dir。保留工作流/应用程序对数据集、规格、输出目录、GPU/平台设置、父检查点和 automl_policy 的覆盖设置。仅在 automl_policy: off 或随包提供的 train 架构/模板缺失时,才直接进行模型训练;如果缺少架构,请说明 AutoML 已启用,但在为此模型生成架构之前无法运行。

evaluate、inference、export 和部署流程等非 train 操作仍由此模型 skill 处理。每次运行的 automl_policy 覆盖设置不会更改模型元数据。

训练要求

  • 数据集类型: segmentation
  • 格式: odvg、coco、coco_raw
  • 监控指标: val_loss
  • AutoML 指标约定: 使用训练期间产生的 val_loss,并将其最小化。独立评估指标是检查点验证 KPI,而非推荐选择目标。
  • 评估检查点指标: 如果检查点会生成类别预测,则使用 [segm] test_mAP50 和 [bbox] test_mAP50。训练时间非常短的冒烟测试检查点可能成功完成评估,但 KPI 字典为空;遇到这种情况,请验证评估操作及结果产物。训练阶段的 AutoML 选择仍以 val_loss 为依据。
各操作的数据集要求
操作Spec 键来源文件列表?
评估dataset.test_data_sourceseval_datasetimage_dir: images.tar.gz, json_file: annotations.json否
评估dataset.test_data_sources.data_typeeval_datasetOD否
推理dataset.infer_data_sourcesinference_datasetimage_dir: images.tar.gz, captions: text prompts否
推理dataset.infer_data_sources.data_typeinference_datasetOD否
量化dataset.train_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json是
量化dataset.val_data_sourceseval_datasetimage_dir: images.tar.gz, json_file: annotations.json否
量化dataset.val_data_sources.data_typeeval_datasetOD否
量化dataset.quant_calibration_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json否
训练dataset.train_data_sourcestrain_datasetsimage_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json是
训练dataset.val_data_sourceseval_datasetimage_dir: images.tar.gz, json_file: annotations.json否
训练dataset.val_data_sources.data_typeeval_datasetOD否
常见 Spec 覆盖项

每项操作都必须覆盖数据源——Agent 必须根据上方「各操作的数据集要求」表构造数据源路径,并将其纳入 spec_overrides。

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

训练(必须指定数据源):

{
    "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.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.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"]},
}

量化(必须指定数据源):

{
    "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.enc_layers / model.dec_layers:训练/AutoML 运行时,两者都保持为 6。掩码头在验证期间会断言有六个解码器输出,因此复制会减少 Transformer 层数的 Grounding DINO 冒烟测试覆盖项会导致立即失败。
  • AutoML 指标说明:训练阶段的 AutoML 使用 metric="val_loss" 和 direction="minimize"。打包的训练循环会记录验证损失标量;训练作业不会输出 [bbox] val_mAP@50。只有当预测结果通过其阈值时,独立评估才会输出 [segm] test_mAP50 和 [bbox] test_mAP50;不要用这两个有条件的测试指标替代用于对 AutoML 建议进行排名的训练损失。
  • 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 键描述默认值
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。建议 GPU 显存为 24GB 或以上。

错误模式

CUDA 显存不足:减小 batch_size。掩码预测会在 Grounding DINO 的基础上增加额外开销。

test_threshold 的部署架构错误:TAO Deploy 使用 evaluate.text_threshold 和 inference.text_threshold。部署规格中不要使用 test_threshold。

部署模型形状不匹配:将转换时使用的 Transformer 和掩码结构字段带入部署的 evaluate/inference 规格中,包括 model.num_queries、 model.num_select、model.max_text_len、model.num_region_queries 和 model.has_mask。这些值必须与用于构建 TensorRT 引擎的 ONNX 模型相匹配。

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,也不要修改生成的运行器脚本来猜测检查点路径。

在 SDK 解析器之外选择 Mask Grounding DINO 检查点时,应与预期的 epoch/step 制品精确匹配,例如 model_epoch_000_step_00049.pth。仅当明确请求 latest 时, mask_gdino_model_latest.pth 这个符号链接才有效。父级 PyTorch mask_grounding_dino CLI 支持 train、evaluate、inference、export 和 quantize;TensorRT 引擎生成、TensorRT 推理和 TensorRT 评估请通过 references/tao-deploy-mask-grounding-dino.md 执行。

部署

评论

支持 Markdown