tao-train-mask-auto-encoder

MAE

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

MAE（Masked Autoencoder，掩码自编码器）用于自监督预训练和微调。它会随机遮蔽图像块并重建这些图像块，以学习视觉表征。支持预训练和微调阶段。

微调时，为预训练的 MAE 权重设置 train.pretrained_model_path。

对于 TAO Deploy TensorRT 操作（gen_trt_engine），请先阅读 references/tao-deploy-mask-auto-encoder.md。部署规范模板位于此技能的 references/ 文件夹中，文件名前缀为 spec_template_deploy_*.yaml。

父级 PyTorch mae CLI 支持 train、evaluate、inference 和 export。请通过部署工作流构建 TensorRT 引擎，而非使用模型技能。

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。对于“turn off AutoML”“disable AutoML”“no HPO”或“plain training”等说法，仅在本次运行中将其视为 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，但在生成架构之前无法运行 AutoML。

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

训练要求

• 数据集类型： image_classification
• 格式： ssl
• 接受的数据集意图： training、evaluation、testing
• 预训练监控指标： train_loss，取最小值。
• AutoML 指标约定： 微调时使用 ACC_all，并采用最大化方向。微调验证和打包的 evaluate 操作会输出 ACC_all（以及 val_loss）；不会输出 train_loss。当 train.stage: finetune 时，使用 ACC_all 选择试验和最终检查点。

每种操作的数据集要求

操作	配置键	来源	文件	是否为列表？
train	dataset.train_data_sources	train_datasets	images_train.tar.gz	否
train	dataset.val_data_sources	eval_dataset	images_val.tar.gz	否
evaluate	dataset.val_data_sources	eval_dataset	images_val.tar.gz	否
inference	dataset.test_data_sources	inference_dataset	images_test.tar.gz	否

对于 SDK/应用作业输入，images_*.tar.gz 压缩包会作为操作输入上传。对于使用挂载到主机的数据直接运行本地 Docker 的情况，请先解压压缩包，并将 dataset.train_data_sources、dataset.val_data_sources 和 dataset.test_data_sources 指向解压后的 images_train、images_val 和 images_test 文件夹。直接将本地 tar 路径传给 MAE CLI 可能会导致数据加载器中样本数为零，因为本地数据加载器不会解包该压缩包路径。

典型配置覆盖项

每种操作都必须覆盖数据源——Agent 必须根据上方“每种操作的数据集要求”表构造数据源路径，并将其包含在 spec_overrides 中。

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

train（必需的数据源）：

{
    "dataset.train_data_sources": f"{S3_TRAIN}/images_train.tar.gz",
    "dataset.val_data_sources": f"{S3_EVAL}/images_val.tar.gz",
    "train.num_epochs": 10,
    "train.optim.lr": 2e-4,
}

evaluate（必需的数据源）：

{
    "dataset.val_data_sources": f"{S3_EVAL}/images_val.tar.gz",
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "train.stage": "finetune",
}

inference（必需的数据源）：

{
    "dataset.test_data_sources": f"{S3_EVAL}/images_test.tar.gz",
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "train.stage": "finetune",
}

评估数据集

可选。预训练不需要评估数据。微调可以选择使用验证集。

重要参数

• train.stage：训练阶段。选项：pretrain、finetune。预训练通过掩码学习表征。微调会添加分类头。
• model.arch：架构。默认值为 convnextv2_base。对于本地冒烟测试 AutoML，请使用 convnextv2_atto，不要使用 vit_tiny_patch16 等不受支持的名称。支持的系列包括 vit_base_patch16 及更大型的 ViT、ConvNeXtV2 的 atto/femto/pico/nano/tiny/base/large/huge，以及 Hiera 的 tiny/small/base/large/huge。
• model.num_classes：微调的类别数。默认值为 1000（ImageNet）。仅在 finetune 阶段相关。
• model.mask_ratio：预训练期间要遮盖的图像块比例。通常为 0.75。
• model.norm_pix_loss：是否在重建损失中对像素值进行归一化。
• dataset.augmentation.input_size：对于 ConvNeXtV2 MAE，将本地冒烟测试配置保持为 224。降低到 112 可能导致 MAE 掩码网格与特征图维度不兼容。
• MAE 不提供 dataset.workers 配置字段。不要将其添加到冒烟测试覆盖项中；Hydra 会在训练前拒绝未知的 dataset 键。
• train.optim.lr：学习率。默认值为 2e-4。
• dataset.augmentation：增强设置，包括用于微调的 mixup、cutmix。

多 GPU / 多节点

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

配置键	说明	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]
train.num_nodes	节点数量	1
train.distributed_strategy	ddp 或 fsdp	ddp

• ddp 使用 find_unused_parameters=True
• fsdp 强制使用 FP16
• 强烈建议使用多 GPU 进行预训练（需要较大的批次大小）

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

硬件

至少需要 2 个 GPU，建议使用 8 个 GPU。每个 GPU 需要 24 GB 以上显存（推荐 A100）。MAE 预训练受益于在多块 GPU 上使用较大的批次大小。微调对资源的要求较低。

错误模式

阶段不匹配：确保 train.stage 与你的意图一致（预训练还是微调）。在没有 pretrained_model_path 的情况下进行微调会从头开始训练。

使用预训练检查点进行推理：MAE 预测数据加载器在 train.stage: pretrain 时会引发 NotImplementedError。推理和分类式评估请使用 finetune 检查点，或将仅预训练的运行限制为训练／评估／导出。

num_classes 不匹配（仅限微调）：微调时，确保 model.num_classes 与数据集的类别数一致。

Spec 参数／父模型推理

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

TAO Core mae.config.json 中的推理映射：

操作	Spec 字段	推理函数	含义
评估	encryption_key	key	加密密钥
评估	evaluate.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
评估	evaluate.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
评估	results_dir	output_dir	当前作业结果目录
导出	encryption_key	key	加密密钥
导出	export.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
导出	export.onnx_file	create_onnx_file	输出 ONNX 路径
导出	results_dir	output_dir	当前作业结果目录
推理	encryption_key	key	加密密钥
推理	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
推理	inference.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
推理	results_dir	output_dir	当前作业结果目录
训练	encryption_key	key	加密密钥
训练	results_dir	output_dir	当前作业结果目录
训练	train.pretrained_model_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
训练	train.resume_training_checkpoint_path	resume_model	从当前作业结果文件夹推断出的模型文件

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

在 SDK 解析器之外解析检查点时，应准确选择目标 epoch／step 产物，例如 model_epoch_000_step_00099.pth。 仅当明确要求使用最新版本时，才使用 convnextv2_atto_latest.pth 或其他最新版本符号链接。将 train.stage、model.arch、model.num_classes 和导出输入尺寸传递到评估、推理、导出和部署规格中，确保检查点与 ONNX／引擎形状相匹配。

部署

• tao-deploy-mask-auto-encoder