tao-train-mask-auto-label

MAL

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

MAL（Mask Auto-Label，掩码自动标注）用于弱监督分割。它可以根据最少的标注（例如点或框标注）生成分割掩码。使用 ViT-MAE 骨干网络。

设置 train.pretrained_model_path，指定 ViT-MAE 预训练权重。

快速开始（docker run）

原生 Docker 启动——主机上无需 TAO SDK 或 Python。如果本地 Docker/平台技能提供了更严格的环境专用命令（非 root UID 映射、缓存重定向、远程守护进程），请改用该命令。

TAO_PYT_IMAGE_DEFAULT=nvcr.io/nvidia/tao/tao-toolkit:7.3.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 接收其配置文件；results_dir 在配置文件中设置，也可以通过命令行覆盖。挂载配置文件引用的任何预训练权重目录，并确保各项操作使用的容器内路径保持一致。

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；维护者须在打包技能库之前重新生成架构和模板。

训练操作策略

此模型在模型层启用了 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。保留工作流/应用针对数据集、配置文件、输出目录、GPU/平台设置、父检查点和 automl_policy 所做的覆盖。只有在 automl_policy: off 或缺少随包提供的训练架构/模板时，才直接训练模型；若缺少架构/模板，请说明 AutoML 虽已启用，但在生成架构之前无法用于此模型。

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

训练要求

• **数据集类型：**分割
• **格式：**默认
• **监控指标：**mIoU
• **AutoML 指标约定：**使用 evaluate 操作输出的 mIoU，并将其最大化。将记录的 AutoML 目标与评估器输出的 mIoU 进行比较；不要替换为训练损失。

各操作的数据集要求

操作	配置键	来源	文件	是否为列表？
evaluate	dataset.val_img_dir	eval_dataset	images.tar.gz	否
evaluate	dataset.val_ann_path	eval_dataset	annotations.json	否
inference	inference.img_dir	inference_dataset	images.tar.gz	否
inference	inference.ann_path	inference_dataset	annotations.json	否
train	dataset.train_img_dir	train_datasets	images.tar.gz	否
train	dataset.train_ann_path	train_datasets	annotations.json	否
train	dataset.val_img_dir	eval_dataset	images.tar.gz	否
train	dataset.val_ann_path	eval_dataset	annotations.json	否

典型配置覆盖项

每项操作都必须覆盖数据源——Agent 必须根据上方「各操作的数据集要求」表构造数据源路径，并将其包含在 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 符号链接。

评估数据集

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

重要参数

• 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：达到完整学习率前的预热轮数。
• dataset.load_mask：标注是否包含预先计算的分割掩码。执行训练或推理且没有掩码真实标注时，如果 COCO 标注包含框但没有 segmentation 字段，请将此项设为 false。如果每条标注都包含分割数据，则保持为 true。按 mIoU 选择结果的 AutoML 验证需要分割真实标注，并要求 dataset.load_mask: true；仅包含边界框的评估会产生非有限的 mIoU，不得将其作为有效的 AutoML 目标。

AutoML / HPO 说明

启动 MAL AutoML 时，请保持默认的试运行搜索空间较窄，并传入 automl_hyperparameters=["train.lr", "train.wd"]。使用围绕 ViT-MAE 微调默认值的保守贝叶斯范围，例如 train.lr 取值范围为 1e-7 至 1e-5，train.wd 取值范围为 1e-5 至 1e-2。 打包的训练 schema 将这两个参数标记为默认 AutoML 参数；如果使用的运行时仍从其捆绑配置模块中推导 MAL 搜索元数据，请显式传入这两个参数。

多 GPU / 多节点

**启动方式：**由 Lightning 管理（单个 python 进程，由 Lightning 启动工作进程）。

配置键	说明	默认值
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。

硬件

至少需要 1 个 GPU，建议使用 2 个 GPU。每个 GPU 需配备 24GB 或更大的显存（推荐 A100）。ViT-MAE 主干网络在 crop_size=512 时需要 24GB 或更大的 GPU 显存。

错误模式

CUDA out of memory：减小 dataset.crop_size（512 -> 384 -> 256），或使用较小的 ViT-MAE 变体（base 与 large）。

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

Spec 参数 / 父模型推断

特定于模型的推断映射应写在此 MD 文件中，而不是 config.json 中。生成的运行器应读取此部分，并在调用 create_job() 前使用 SDK 辅助函数应用这些映射。这与旧版微服务的 infer_params.py 流程一致。

TAO Core mal.config.json 中的推断映射：

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

对于 parent_model 或 parent_model_folder，请将上游 train/export/AutoML 子作业 ID 作为 parent_job_id 传入。SDK 会列出父作业结果文件夹、筛选检查点制品，并返回选定的模型文件或文件夹。请勿将这些映射重新添加到 config.json 中，也不要修改生成的运行器脚本来猜测检查点路径。