tao-train-mask-auto-label

MAL (Mask Auto-Label) for weakly-supervised segmentation. Produces segmentation masks from minimal annotations (point or box annotations) using a ViT-MAE backbone. Use when training, evaluating, or running inference for a TAO MAL model. Trigger phrases include "train MAL", "Mask Auto-Label", "weakly

By nvidia · 1,549 installs

npx skills add nvidia/skills --skill tao-train-mask-auto-label

Source repository · Upstream listing

MAL 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). MAL (Mask Auto Label) for weakly supervised segmentation. Produces segmentation masks from minimal annotations (e.g., point or box annotations). Uses ViT MAE backbone. Set train.pretrained model path for ViT MAE pretrained weights. Quick Start (docker run) Docker native launch — no TAO SDK and no Python on the host. Use the local Docker/platform skill instead when it gives a stricter environment specific command (non root UID mapping, cache redirects, remote daemons). Train: Evaluate: Inference: Every action takes its spec with e ; results dir is set in the spec or overridden on the command line. Mount any pretrained weights directory the spec references, and keep every in container path consistent across actions. 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: default Monitoring metric: mIoU Per Action Dataset Requirements Action Spec Key Source Files List? evaluate dataset.val img dir eval dataset images.tar.gz No evaluate dataset.val ann path eval dataset annotations.json No inference inference.img dir inference dataset images.tar.gz No inference inference.ann path inference dataset annotations.json No train dataset.train img dir train datasets images.tar.gz No train dataset.train ann path train datasets annotations.json No train dataset.val img dir eval dataset images.tar.gz No train dataset.val ann path eval dataset annotations.json 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 . MAL expects COCO style annotation JSON plus image paths that match the JSON file name entries after the data source is prepared. Archive only CSV/image datasets are not compatible unless they are converted to this format first. train (mandatory data sources): evaluate (mandatory data sources): inference (mandatory data sources): For checkpoint dependent actions, use the model resolver declared in references/skill info.yaml . Select the exact epoch/step checkpoint requested by the user or the best checkpoint when a best checkpoint action is requested. The mal model latest.pth symlink is only appropriate when the user explicitly asks for the latest checkpoint. Eval Dataset Optional. Val images and annotations configured alongside train paths. Important Parameters model.arch : ViT MAE backbone variant. Default vit mae base/16. Avoid vit deit tiny/16 ; the current runtime rejects tiny ViT variants. train.lr : Learning rate. Default 1e 6 (very low — fine tuning ViT). dataset.crop size : Training crop size. Default 512. Use this key, not model.crop size . train.warmup epochs : Warmup epochs before full learning rate. model.load mask : Whether to load pre computed masks. AutoML / HPO Notes For MAL AutoML launches, keep the default smoke search space narrow and pass automl hyperparameters=["train.lr", "train.wd"] . Use conservative Bayesian ranges around the ViT MAE fine tuning defaults, for example train.lr from 1e 7 to 1e 5 and train.wd from 1e 5 to 1e 2 . The packaged train schema marks these two parameters as the default AutoML parameters; pass them explicitly when using a runtime that still derives MAL search metadata from its bundled config module. Multi GPU / Multi Node Launch method: Lightning managed (single python process, Lightning spawns workers). 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 Multi GPU strategy: ddp find unused parameters true No fsdp support LR auto scaling: lr = lr num devices batch size (learning rate is scaled automatically by device count and batch size) Multi node env vars (set by orchestrator): WORLD SIZE , NODE RANK , MASTER ADDR , MASTER PORT , NUM GPU PER NODE . Hardware Minimum 1 GPU(s), recommended 2 GPU(s). 24GB+ (A100 recommended) VRAM per GPU. ViT MAE backbone at crop size=512 needs 24GB+ GPU memory. Error Patterns CUDA out of memory : Reduce dataset.crop size (512 384 256) or use a smaller ViT MAE variant (base vs large). Key crop size not in MALModelConfig : The crop size override was placed under model.crop size . Move it to dataset.crop size . 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 mal.config.json : Action Spec Field Inference Function Meaning evaluate evaluate.checkpoint parent model model file inferred from the parent job results folder evaluate results dir output dir current job results directory inference inference.checkpoint parent model model file inferred from the parent job results folder inference inference.label dump path create inference result file mal MAL inference JSON path inference results dir output dir current job results directory train train.pretrained model path ptm if no resume model optional pretrained model when not resuming train train.resume training checkpoint path resume model exact checkpoint for resume runs train results dir output dir current job results directory 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.