tao-train-mask-grounding-dino

name	tao-train-mask-grounding-dino
description	Mask Grounding DINO for grounded instance segmentation. Extends Grounding DINO with a mask-prediction head for open-set segmentation guided by text prompts. Use when training, evaluating, exporting, quantizing, or running inference for a TAO Mask-Grounding-DINO model. Trigger phrases include "train Mask Grounding DINO", "open-vocabulary segmentation", "text-prompted instance segmentation", "grounded mask DETR".
license	Apache-2.0
compatibility	Requires docker + nvidia-container-toolkit.
metadata	{"version":"0.1.0","author":"NVIDIA Corporation"}
allowed-tools	Read Bash
tags	["segmentation"]

Mask Grounding DINO

Standalone install? If this session was not initialized by the TAO skill bank plugin, run the tao-setup skill first (host preflight, credentials, cross-skill discovery).

Mask Grounding DINO for grounded instance segmentation. Extends Grounding DINO with mask prediction head for open-set segmentation guided by text prompts.

Set train.pretrained_model_path for full model weights.

For TAO Deploy TensorRT actions (gen_trt_engine, TensorRT evaluate, and TensorRT inference), read references/tao-deploy-mask-grounding-dino.md first. Deploy spec templates live in this skill's references/ folder with the spec_template_deploy_*.yaml prefix.

Dataclass Schemas

Generated TAO Core schemas are packaged in schemas/<action>.schema.json, with schemas/manifest.json listing available actions. Each generated schema also emits references/spec_template_<action>.yaml from the schema top-level default field. AutoML enablement is declared at the model layer in references/skill_info.yaml via automl_enabled. Runnable AutoML for an action requires schemas/<action>.schema.json and references/spec_template_<action>.yaml to exist and parse. Use the packaged selected-action schema for automl_default_parameters, automl_disabled_parameters, defaults, min/max bounds, enums, option weights, math conditions, dependencies, and popular parameters. Do not expect ~/tao-core at runtime; maintainers regenerate schemas/templates before packaging the skill bank.

Train Action Policy

This model is AutoML-enabled at the model layer. Before handling any train-stage request, read references/skill_info.yaml and resolve the run override from either an explicit automl_policy value or the user's workflow request. Use automl_policy: on by default and only expose on / off in new launch prompts. Treat phrases like "turn off AutoML", "disable AutoML", "no HPO", or "plain training" as automl_policy: off for this run only. When automl_policy: on, automl_enabled: true, and both schemas/train.schema.json and references/spec_template_train.yaml are packaged, route the train action through tao-skill-bank:tao-run-automl by default with this model's skill_dir. Preserve workflow/application overrides for datasets, specs, output directories, GPU/platform settings, parent checkpoints, and automl_policy. Use direct model training only when automl_policy: off or the packaged train schema/template is missing; in the missing-schema case, report that AutoML is enabled but not runnable for this model until schemas are generated.

Non-train actions such as evaluate, inference, export, and deploy flows stay in this model skill. The per-run automl_policy override does not change model metadata.

Training Requirements

• Dataset type: segmentation
• Formats: odvg, coco, coco_raw
• Monitoring metric: val_loss
• AutoML metric contract: Use val_loss emitted during training and minimize it. Standalone evaluation metrics are checkpoint-validation KPIs, not the recommendation-selection objective.
• Evaluation checkpoint metrics: [segm] test_mAP50 and [bbox] test_mAP50 when the checkpoint produces class predictions. A very short smoke-trained checkpoint can complete evaluation successfully with an empty KPI dictionary; verify the evaluation action and result artifact in that case. Train-stage AutoML selection remains based on val_loss.

Per-Action Dataset Requirements

Action	Spec Key	Source	Files	List?
evaluate	dataset.test_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	No
evaluate	dataset.test_data_sources.data_type	eval_dataset	OD	No
inference	dataset.infer_data_sources	inference_dataset	image_dir: images.tar.gz, captions: text prompts	No
inference	dataset.infer_data_sources.data_type	inference_dataset	OD	No
quantize	dataset.train_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json	Yes
quantize	dataset.val_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	No
quantize	dataset.val_data_sources.data_type	eval_dataset	OD	No
quantize	dataset.quant_calibration_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json	No
train	dataset.train_data_sources	train_datasets	image_dir: images.tar.gz, json_file: annotations_odvg.jsonl, label_map: annotations_odvg_labelmap.json	Yes
train	dataset.val_data_sources	eval_dataset	image_dir: images.tar.gz, json_file: annotations.json	No
train	dataset.val_data_sources.data_type	eval_dataset	OD	No

Typical Spec Overrides

Data source overrides are mandatory for every action — the agent MUST construct data source paths from the Per-Action Dataset Requirements table above and include them in spec_overrides.

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

train (mandatory data sources):

{
    "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 (mandatory data sources):

{
    "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 (mandatory data sources):

{
    "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 (mandatory data sources):

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

Eval Dataset

Optional. Validation uses COCO-format annotations even when training uses ODVG.

Important Parameters

• model.backbone: Default swin_tiny_224_1k. Same backbone options as Grounding DINO.
• train.optim.lr: Learning rate. Default 2e-4. lr_backbone 2e-5. Reuses GDINOTrainExpConfig — same training setup as Grounding DINO.
• model.num_queries: Object queries. Default 900.
• model.enc_layers / model.dec_layers: Keep both at 6 for train/AutoML runs. The mask head asserts six decoder outputs during validation, so copying Grounding DINO smoke overrides that reduce transformer layers causes an immediate failure.
• AutoML metric note: Use metric="val_loss" with direction="minimize" for train-stage AutoML. The packaged train loop logs validation loss scalars; it does not emit [bbox] val_mAP@50 during the train job. Standalone evaluation emits [segm] test_mAP50 and [bbox] test_mAP50 only when predictions survive its thresholds; do not substitute either conditional test metric for the training loss used to rank AutoML recommendations.
• model.has_mask: Enables mask prediction head. Default True. Adds mask/dice/rela loss coefficients.
• model.num_region_queries: Number of region queries for mask prediction. Default 100.
• model.loss_types: Loss components. Default [labels, boxes, masks]. Includes mask_loss_coef, dice_loss_coef, rela_loss_coef.
• evaluate.ioi_threshold: IoI threshold for mask evaluation. Default 0.5.
• evaluate.nms_threshold: NMS threshold. Default 0.2.
• evaluate.text_threshold: Text matching threshold. Default 0.3.
• dataset.has_mask: Dataset includes mask annotations. Default True. val_data_sources default data_type is "VG".

Multi-GPU / Multi-Node

Launch method: Lightning-managed. Same DDP/FSDP behavior as Grounding DINO.

Spec Key	Description	Default
train.num_gpus	Number of GPUs	1
train.gpu_ids	GPU device indices	[0]
train.num_nodes	Number of nodes	1
train.distributed_strategy	ddp or fsdp	ddp

Hardware

Minimum 1 GPU(s), recommended 4 GPU(s). 24GB+ (A100 recommended) VRAM per GPU. Heavier than Grounding DINO due to mask prediction head. 24GB+ GPU memory recommended.

Error Patterns

CUDA out of memory: Reduce batch_size. Mask prediction adds overhead on top of Grounding DINO.

Deploy schema error for test_threshold: TAO Deploy uses evaluate.text_threshold and inference.text_threshold. Do not use test_threshold in deploy specs.

Deploy model shape mismatch: Carry transformer and mask structure fields from export into deploy evaluate/inference specs, including model.num_queries, model.num_select, model.max_text_len, model.num_region_queries, and model.has_mask. These values must match the ONNX model used to build the TensorRT engine.

Spec Param / Parent Model Inference

Model-specific inference mappings belong in this MD file, not in config.json. Generated runners should read this section and apply the mappings with SDK helpers before create_job(). This mirrors the old microservices infer_params.py flow.

Inference mappings from TAO Core mask_grounding_dino.config.json:

Action	Spec Field	Inference Function	Meaning
evaluate	encryption_key	key	encryption key
evaluate	evaluate.checkpoint	parent_model	model file inferred from the parent job results folder
evaluate	evaluate.trt_engine	parent_model	model file inferred from the parent job results folder
evaluate	results_dir	output_dir	current job results directory
export	encryption_key	key	encryption key
export	export.checkpoint	parent_model	model file inferred from the parent job results folder
export	export.onnx_file	create_onnx_file	output ONNX path
export	results_dir	output_dir	current job results directory
gen_trt_engine	encryption_key	key	encryption key
gen_trt_engine	gen_trt_engine.onnx_file	parent_model	model file inferred from the parent job results folder
gen_trt_engine	gen_trt_engine.trt_engine	create_engine_file	output TensorRT engine path
gen_trt_engine	results_dir	output_dir	current job results directory
inference	encryption_key	key	encryption key
inference	inference.checkpoint

For parent_model or parent_model_folder, pass the upstream train/export/AutoML child job id as parent_job_id. The SDK lists the parent result folder, filters checkpoint artifacts, and returns the selected model file or folder. Do not add these mappings back to config.json and do not patch generated runner scripts to guess checkpoint paths.

When selecting a Mask Grounding DINO checkpoint outside the SDK resolver, match the intended epoch/step artifact exactly, for example model_epoch_000_step_00049.pth. The mask_gdino_model_latest.pth symlink is valid only when latest is explicitly requested. The parent PyTorch mask_grounding_dino CLI supports train, evaluate, inference, export, and quantize; run TensorRT engine generation, TensorRT inference, and TensorRT evaluation through references/tao-deploy-mask-grounding-dino.md.

Deployment

• tao-deploy-mask-grounding-dino