⚡

掩码自动标注（MAL）训练Skill tao-train-mask-auto-label

该技能用于训练、评估和推理基于ViT-MAE骨干的MAL（掩码自动标注）弱监督分割模型。它能够从点或框等最小标注中自动生成分割掩码，支持COCO格式数据、AutoML超参数搜索、多GPU/多节点训练，并输出mIoU指标。适用于视觉AI中的分割模型开发。关键词：弱监督分割、掩码自动标注、MAL、ViT-MAE、分割模型训练、最小标注、AutoML、COCO格式、视觉AI、TAO。

视觉模型训练 0 次安装 0 次浏览 更新于 9/6/2026
名称 tao-train-mask-auto-label
描述 MAL（掩码自动标注）用于弱监督分割。使用ViT-MAE骨干网络，从最小化标注（点或框标注）生成分割掩码。当需要对TAO MAL模型进行训练、评估或推理时使用。触发短语包括“train MAL”、“Mask Auto-Label”、“weakly-supervised segmentation”、“box-prompted segmentation”、“minimal-annotation mask prediction”。
开源协议 Apache-2.0 compatibility: Requires docker + nvidia-container-toolkit. metadata:
版本 “0.1.0”
作者 NVIDIA Corporation allowed-tools: Read Bash tags: - segmentation

MAL

独立安装？ 如果此会话不是由TAO技能库插件初始化的，请先运行tao-setup技能（主机预检、凭据、跨技能发现）。

MAL（掩码自动标注）用于弱监督分割。从最小标注（例如点或框标注）生成分割掩码。使用ViT-MAE骨干网络。

为ViT-MAE预训练权重设置train.pretrained_model_path。

Quick Start (docker run) 快速开始（docker run）

Docker本地启动——不需要TAO SDK，也不需要主机上的Python。当本地Docker/平台技能给出更严格的环境特定命令（非root UID映射、缓存重定向、远程守护进程）时，请改为使用该技能。

TAO_PYT_IMAGE_DEFAULT=nvcr.io/nvidia/tao/tao-toolkit:7.1.0-pyt  # versions-key: images.tao_toolkit.pyt
TAO_PYT_IMAGE="${TAO_PYT_IMAGE:-$TAO_PYT_IMAGE_DEFAULT}"
RUN_ROOT="${RUN_ROOT:-$PWD}"
DOCKER_COMMON=(
  --rm --gpus all --shm-size=8g
  --shm-size=8g
  --ulimit memlock=-1
  --ulimit stack=67108864
  -v "$RUN_ROOT/data:/data:ro"
  -v "$RUN_ROOT/specs:/specs:ro"
  -v "$RUN_ROOT/results:/results"
)

训练：

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  mal train -e /specs/train.yaml

评估：

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  mal evaluate -e /specs/evaluate.yaml

推理：

docker run "${DOCKER_COMMON[@]}" "$TAO_PYT_IMAGE" \
  mal inference -e /specs/inference.yaml

每个操作都通过-e参数接受其spec；results_dir在spec中设置，或在命令行中覆盖。挂载spec引用的任何预训练权重目录，并保持所有操作中容器内路径的一致性。

Dataclass Schemas 数据类模式

生成的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；维护者在打包技能库前会重新生成模式/模板。

Train Action Policy 训练操作策略

该模型在模型层启用了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使用此模型的skill_dir路由训练操作。保留工作流/应用程序对数据集、spec、输出目录、GPU/平台设置、父检查点和automl_policy的覆盖。仅在automl_policy: off或打包的训练模式/模板缺失时，才直接进行模型训练；在模式缺失的情况下，报告AutoML已启用但在生成模式之前无法运行。

非训练操作，如evaluate、inference、export和部署流程，仍保留在此模型技能中。每次运行的automl_policy覆盖不会改变模型元数据。

Training Requirements 训练要求

  • 数据集类型： 分割
  • 格式： 默认
  • 监控指标： mIoU

Per-Action Dataset Requirements 每种操作的数据集要求

操作 Spec键 来源 文件 列表?
评估 dataset.val_img_dir eval_dataset images.tar.gz 否
评估 dataset.val_ann_path eval_dataset annotations.json 否
推理 inference.img_dir inference_dataset images.tar.gz 否
推理 inference.ann_path inference_dataset annotations.json 否
训练 dataset.train_img_dir train_datasets images.tar.gz 否
训练 dataset.train_ann_path train_datasets annotations.json 否
训练 dataset.val_img_dir eval_dataset images.tar.gz 否
训练 dataset.val_ann_path eval_dataset annotations.json 否

Typical Spec Overrides 典型Spec覆盖

数据源覆盖对每个操作都是强制性的——代理必须根据上面的每种操作数据集要求表构造数据源路径，并将它们包含在spec_overrides中。 MAL期望COCO风格的标注JSON以及图像路径，这些路径在数据源准备好后与JSON的file_name条目匹配。仅包含存档的CSV/图像数据集不兼容，除非它们先转换为这种格式。

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

train（强制数据源）：

{
    "train.num_gpus": 1,
    "train.gpu_ids": [
        0
    ],
    "train.num_epochs": 5,
    "train.checkpoint_interval": 5,
    "train.validation_interval": 5,
    "dataset.train_img_dir": f"{S3_TRAIN}/images.tar.gz",
    "dataset.train_ann_path": f"{S3_TRAIN}/annotations.json",
    "dataset.val_img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val_ann_path": f"{S3_EVAL}/annotations.json",
}

evaluate（强制数据源）：

{
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.val_img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val_ann_path": f"{S3_EVAL}/annotations.json",
}

inference（强制数据源）：

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "inference.img_dir": f"{S3_EVAL}/images.tar.gz",
    "inference.ann_path": f"{S3_EVAL}/annotations.json",
}

对于依赖检查点的操作，请使用references/skill_info.yaml中声明的模型解析器。当用户请求特定epoch/step检查点时选择该检查点，或者当请求最佳检查点操作时选择最佳检查点。mal_model_latest.pth符号链接仅在用户明确要求最新检查点时使用。

Eval Dataset 评估数据集

可选。与训练路径一起配置验证图像和标注。

Important Parameters 重要参数

  • model.arch：ViT-MAE骨干网络变体。默认vit-mae-base/16。避免使用vit-deit-tiny/16；当前运行时拒绝tiny ViT变体。
  • train.lr：学习率。默认1e-6（非常低——微调ViT）。
  • dataset.crop_size：训练裁剪尺寸。默认512。使用此键，而不是model.crop_size。
  • train.warmup_epochs：达到完整学习率前的预热epoch数。
  • model.load_mask：是否加载预计算的掩码。

AutoML / HPO Notes AutoML/HPO说明

对于MAL AutoML启动，将默认冒烟搜索空间保持较窄，并传递automl_hyperparameters=["train.lr", "train.wd"]。使用ViT-MAE微调默认值附近的保守贝叶斯范围，例如train.lr从1e-7到1e-5，train.wd从1e-5到1e-2。打包的训练模式将这两个参数标记为默认AutoML参数；当使用仍从其捆绑配置模块派生MAL搜索元数据的运行时，请显式传递它们。

Multi-GPU / Multi-Node 多GPU/多节点

启动方式： Lightning管理（单个python进程，Lightning生成worker）。

Spec键 描述 默认值
train.num_gpus GPU数量 1
train.gpu_ids GPU设备索引 [0]
train.num_nodes 节点数量 1
  • 多GPU策略：ddp_find_unused_parameters_true
  • 不支持fsdp
  • 学习率自动缩放： lr = lr * num_devices * batch_size（学习率根据设备数和批量大小自动缩放）

多节点环境变量（由编排器设置）：WORLD_SIZE、NODE_RANK、MASTER_ADDR、MASTER_PORT、NUM_GPU_PER_NODE。

Hardware 硬件

最少1个GPU，推荐2个GPU。每个GPU需要24GB以上（推荐A100）显存。在crop_size=512时，ViT-MAE骨干网络需要24GB以上GPU内存。

Error Patterns 错误模式

CUDA内存不足：减小dataset.crop_size（512 -> 384 -> 256）或使用更小的ViT-MAE变体（base对large）。

MALModelConfig中不存在键crop_size：裁剪尺寸覆盖被放在了model.crop_size下。请将其移动到dataset.crop_size。

Spec Param / Parent Model Inference Spec参数/父模型推断

模型特定的推理映射应放在此MD文件中，而不是config.json中。生成的运行器应在create_job()之前使用SDK辅助程序读取此部分并应用映射。这类似于旧的微服务infer_params.py流程。

来自TAO Core mal.config.json的推理映射：

操作 Spec字段 推理函数 含义
评估 evaluate.checkpoint parent_model 从父作业结果文件夹推断出的模型文件
评估 results_dir output_dir 当前作业结果目录
推理 inference.checkpoint parent_model 从父作业结果文件夹推断出的模型文件
推理 inference.label_dump_path create_inference_result_file_mal MAL推理JSON路径
推理 results_dir output_dir 当前作业结果目录
训练 train.pretrained_model_path ptm_if_no_resume_model 不恢复时可选的预训练模型
训练 train.resume_training_checkpoint_path resume_model 恢复运行的确切检查点
训练 results_dir output_dir 当前作业结果目录

对于parent_model或parent_model_folder，将上游训练/导出/AutoML子作业ID作为parent_job_id传递。SDK列出父结果文件夹，过滤检查点工件，并返回所选模型文件或文件夹。不要将这些映射添加回config.json，也不要修补生成的运行器脚本来猜测检查点路径。