tao-train-mask-auto-label

MAL

面向弱监督分割的 MAL（Mask Auto-Label）。通过少量标注（例如点或框标注）生成分割掩码。使用 ViT-MAE 主干网络。

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

数据类架构

生成的 TAO Core 架构打包在 schemas/<action>.schema.json 中，schemas/manifest.json 列出了可用操作。每个生成的架构还会根据架构顶层的 default 字段生成 references/spec_template_<action>.yaml。AutoML 是否启用由模型层通过 references/skill_info.yaml 中的 automl_enabled 声明。要运行 AutoML，仍需存在且可解析的 schemas/train.schema.json 和 references/spec_template_train.yaml。使用打包的训练架构获取 automl_default_parameters、automl_disabled_parameters、默认值、最小值／最大值范围、枚举、选项权重、数学条件、依赖项和常用参数。不要期望运行时存在 ~/tao-core；维护者会在打包技能库之前重新生成架构／模板。

训练操作策略

此模型在模型层启用了 AutoML。在处理任何训练阶段请求之前，请读取 references/skill_info.yaml，并根据显式的 automl_policy 值或用户的工作流请求确定本次运行的策略。将“关闭 AutoML”“禁用 AutoML”“不使用 HPO”或“普通训练”等说法视为仅对本次运行设置 automl_policy: off；否则默认使用 auto。当 automl_policy: auto、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，但在生成架构之前无法运行 AutoML。

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

训练要求

• 数据集类型： 分割
• 格式： 默认
• 监控指标： 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 中。

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（必需的数据源）：

{
    "dataset.val_img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val_ann_path": f"{S3_EVAL}/annotations.json",
}

inference（必需的数据源）：

{
    "inference.img_dir": f"{S3_EVAL}/images.tar.gz",
    "inference.ann_path": f"{S3_EVAL}/annotations.json",
}

评估数据集

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

重要参数

• model.arch：ViT-MAE 主干网络变体。默认值为 vit-mae-base/16。选项包括 vit-mae-large/16 和其他 ViT-MAE 变体。
• train.lr：学习率。默认值为 1e-6（非常低——用于微调 ViT）。
• model.crop_size：训练裁剪尺寸。默认值为 512。
• train.warmup_epochs：达到完整学习率前的预热轮数。
• model.load_mask：是否加载预计算的掩码。

多 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）。当 crop_size=512 时，ViT-MAE 主干网络需要 24GB 或更多 GPU 显存。

错误模式

CUDA 显存不足：降低 model.crop_size（512 -> 384 -> 256），或使用较小的 ViT-MAE 变体（base 而非 large）。

规范参数 / 父模型推断

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

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

操作	规范字段	推断函数	含义
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	results_dir	output_dir	当前作业的结果目录

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