tao-train-mask2former

Mask2Former

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

Mask2Former 用于通用图像分割（全景、实例和语义分割）。它基于 Transformer，并采用掩码注意力机制，可实现高质量的分割结果。

为 Swin 主干网络权重设置 model.backbone.pretrained_weights。

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

数据类架构

生成的 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；维护者会在打包 skill bank 前重新生成架构／模板。

训练操作策略

此模型在模型层启用了 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 覆盖值不会更改模型元数据。

训练要求

• 数据集类型： 分割
• 格式： coco_panoptic、coco
• 监控指标： mIoU
• AutoML 指标： 最大化输出的验证／评估 mIoU；绝不以训练损失替代。

各操作的数据集要求

Action	Spec Key	Source	Files	List?
evaluate	dataset.train.type	train_datasets	coco_panoptic	No
evaluate	dataset.val.type	eval_dataset	coco_panoptic	No
evaluate	dataset.test.type	eval_dataset	coco_panoptic	No
evaluate	dataset.train.img_dir	train_datasets	images.tar.gz	No
evaluate	dataset.label_map	train_datasets	coco_panoptic: label_map_panoptic.json; *: label_map.json	No
evaluate	dataset.train.instance_json	train_datasets	annotations.json	No
evaluate	dataset.train.panoptic_json	train_datasets	annotations_panoptic.json	No
evaluate	dataset.train.panoptic_dir	train_datasets	images_panoptic.tar.gz	No
evaluate	dataset.val.img_dir	eval_dataset	images.tar.gz	No
evaluate	dataset.val.instance_json	eval_dataset	annotations.json	No
evaluate	dataset.val.panoptic_json	eval_dataset	annotations_panoptic.json	No
evaluate	dataset.val.panoptic_dir	eval_dataset	images_panoptic.tar.gz	No
evaluate	dataset.test.img_dir	eval_dataset	images.tar.gz	No
inference	dataset.train.type	train_datasets	coco_panoptic	No
inference	dataset.val.type	eval_dataset	coco_panoptic	No
inference	dataset.test.type	eval_dataset	coco_panoptic	No
inference	dataset.train.img_dir	train_datasets	images.tar.gz	No
inference	dataset.label_map	train_datasets	coco_panoptic: label_map_panoptic.json; *: label_map.json	No
inference	dataset.train.instance_json	train_datasets	annotations.json	No
inference	dataset.train.panoptic_json	train_datasets	annotations_panoptic.json	No
inference	dataset.train.panoptic_dir	train_datasets	images_panoptic.tar.gz	No
inference	dataset.val.img_dir	eval_dataset	images.tar.gz	No
inference	dataset.val.instance_json	eval_dataset	annotations.json	No
inference	dataset.val.panoptic_json	eval_dataset	annotations_panoptic.json	No
inference	dataset.val.panoptic_dir	eval_dataset	images_panoptic.tar.gz	No
inference	dataset.test.img_dir	eval_dataset	images.tar.gz	No
quantize	dataset.train.type	train_datasets	coco_panoptic	No
quantize	dataset.val.type	eval_dataset	coco_panoptic	No
quantize	dataset.test.type	eval_dataset	coco_panoptic	No
quantize	dataset.train.img_dir	train_datasets	images.tar.gz	No
quantize	dataset.label_map	train_datasets	coco_panoptic: label_map_panoptic.json; *: label_map.json	No
quantize	dataset.train.instance_json	train_datasets	annotations.json	No
quantize	dataset.train.panoptic_json	train_datasets	annotations_panoptic.json	No
quantize	dataset.train.panoptic_dir	train_datasets	images_panoptic.tar.gz	No
quantize	dataset.val.img_dir	eval_dataset	images.tar.gz	No
quantize	dataset.val.instance_json	eval_dataset	annotations.json	No
quantize	dataset.val.panoptic_json	eval_dataset	annotations_panoptic.json	No
quantize	dataset.val.panoptic_dir	eval_dataset	images_panoptic.tar.gz	No
quantize	dataset.test.img_dir	eval_dataset	images.tar.gz	No
quantize	dataset.quant_calibration_dataset.images_dir	train_datasets	images.tar.gz	No
train	dataset.train.type	train_datasets	coco_panoptic	No
train	dataset.val.type	eval_dataset	coco_panoptic	No
train	dataset.test.type	eval_dataset	coco_panoptic	No
train	dataset.train.img_dir	train_datasets	images.tar.gz	No
train	dataset.label_map	train_datasets	coco_panoptic: label_map_panoptic.json; *: label_map.json	No
train	dataset.train.instance_json	train_datasets	annotations.json	No
train	dataset.train.panoptic_json	train_datasets	annotations_panoptic.json	No
train	dataset.train.panoptic_dir	train_datasets	images_panoptic.tar.gz	No
train	dataset.val.img_dir	eval_dataset	images.tar.gz	No
train	dataset.val.instance_json	eval_dataset	annotations.json	No
train	dataset.val.panoptic_json	eval_dataset	annotations_panoptic.json	No
train	dataset.val.panoptic_dir	eval_dataset	images_panoptic.tar.gz	No
train	dataset.test.img_dir	eval_dataset	images.tar.gz	No

常见规范覆盖项

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

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

train（必需数据源）：

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "model.sem_seg_head.num_classes": 133,
    "dataset.contiguous_id": True,
    "dataset.train.type": "coco_panoptic",
    "dataset.val.type": "coco_panoptic",
    "dataset.test.type": "coco_panoptic",
    "dataset.train.img_dir": f"{S3_TRAIN}/images.tar.gz",
    "dataset.label_map": f"{S3_TRAIN}/label_map_panoptic.json",
    "dataset.train.instance_json": f"{S3_TRAIN}/annotations.json",
    "dataset.train.panoptic_json": f"{S3_TRAIN}/annotations_panoptic.json",
    "dataset.train.panoptic_dir": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.instance_json": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic_json": f"{S3_EVAL}/annotations_panoptic.json",
    "dataset.val.panoptic_dir": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.img_dir": f"{S3_EVAL}/images.tar.gz",
}

evaluate（必需数据源）：

{
    "evaluate.checkpoint": "<selected train/AutoML checkpoint>",
    "model.sem_seg_head.num_classes": 133,
    "dataset.contiguous_id": True,
    "dataset.train.type": "coco_panoptic",
    "dataset.val.type": "coco_panoptic",
    "dataset.test.type": "coco_panoptic",
    "dataset.train.img_dir": f"{S3_TRAIN}/images.tar.gz",
    "dataset.label_map": f"{S3_TRAIN}/label_map_panoptic.json",
    "dataset.train.instance_json": f"{S3_TRAIN}/annotations.json",
    "dataset.train.panoptic_json": f"{S3_TRAIN}/annotations_panoptic.json",
    "dataset.train.panoptic_dir": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.instance_json": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic_json": f"{S3_EVAL}/annotations_panoptic.json",
    "dataset.val.panoptic_dir": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.img_dir": f"{S3_EVAL}/images.tar.gz",
}

export：

{
    "export.checkpoint": "<selected train/AutoML checkpoint>",
    "export.onnx_file": "<output ONNX path>",
    "model.sem_seg_head.num_classes": "<same value used for train>",
}

inference（必需数据源）：

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "model.sem_seg_head.num_classes": "<same value used for train>",
    "dataset.contiguous_id": True,
    "dataset.train.img_dir": f"{S3_TRAIN}/images.tar.gz",
    "dataset.label_map": f"{S3_TRAIN}/label_map_panoptic.json",
    "dataset.train.instance_json": f"{S3_TRAIN}/annotations.json",
    "dataset.train.panoptic_json": f"{S3_TRAIN}/annotations_panoptic.json",
    "dataset.train.panoptic_dir": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.instance_json": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic_json": f"{S3_EVAL}/annotations_panoptic.json",
    "dataset.val.panoptic_dir": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.img_dir": f"{S3_EVAL}/images.tar.gz",
}

quantize（必需数据源）：

{
    "quantize.model_path": "<selected train/export artifact>",
    "dataset.train.img_dir": f"{S3_TRAIN}/images.tar.gz",
    "dataset.label_map": f"{S3_TRAIN}/label_map_panoptic.json",
    "dataset.train.instance_json": f"{S3_TRAIN}/annotations.json",
    "dataset.train.panoptic_json": f"{S3_TRAIN}/annotations_panoptic.json",
    "dataset.train.panoptic_dir": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.instance_json": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic_json": f"{S3_EVAL}/annotations_panoptic.json",
    "dataset.val.panoptic_dir": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.img_dir": f"{S3_EVAL}/images.tar.gz",
    "dataset.quant_calibration_dataset.images_dir": f"{S3_TRAIN}/images.tar.gz",
}

评估数据集

可选。验证数据源与训练数据源一起属于数据集配置的一部分。

重要参数

• model.sem_seg_head.num_classes：分割类别数。默认值为 200。必须与标注类别一致。
• model.backbone.swin.type：Swin Transformer 变体。默认值为 tiny。可选值包括 tiny、small、base、large。
• model.mode：分割模式。默认值为 panoptic。可选值：panoptic、instance、semantic。
• train.optim.lr：学习率。默认值为 2e-4（AdamW）。
• dataset.train.batch_size：每个 GPU 的批次大小。默认值为 1。由于 Mask2Former 会进行逐像素预测，因此显存占用较高。
• dataset.contiguous_id：若为 true，将 model.sem_seg_head.num_classes 设置为标签映射类别数。若为 false，将 model.sem_seg_head.num_classes 设置为大于最大原始类别 ID 的值，并在 evaluate、inference、export、deploy 和 quantize 中保持相同设置。COCO panoptic S3 样例有 133 个类别，原始 ID 最大为 200，因此原始 ID 验证使用 num_classes: 201。

多 GPU / 多节点

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

规范键	描述	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]
train.num_nodes	节点数量	1
train.distributed_strategy	ddp 或 fsdp	ddp

• 与 DINO 的 DDP/FSDP 行为相同（支持激活检查点）
• FAN 主干网络会自动启用 sync_batchnorm
• fsdp 会强制使用 FP16

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

导出 / TRT 默认设置

• TRT 数据类型：仅支持 FP32、FP16——不支持 INT8
• 父级 PyTorch mask2former CLI 支持 train、evaluate、 inference、export 和 quantize；请通过 references/tao-deploy-mask2former.md 运行 TensorRT 引擎生成、 TensorRT 推理和 TensorRT 评估。 验证 TensorRT 评估时，请导出语义 ONNX（model.mode: semantic）， 因为当前部署评估器接受语义引擎。
• 导出输入尺寸须与部署模板兼容。打包的默认值 export.input_width: 960 和 export.input_height: 544 可以成功导出并构建 TensorRT 引擎；若将导出尺寸缩小为仅用于验证的尺寸（如 128x128），可能会在生成 ONNX 之前触发 PyTorch ONNX minus_one_pos != -1 形状推断断言。

硬件

至少需要 1 张 GPU，建议使用 4 张 GPU。每张 GPU 需要 24GB 以上显存（建议使用 A100）。Mask2Former 显存占用较高。默认将 batch_size 设为 1 是有原因的。建议使用多 GPU，以获得合理的训练速度。

错误模式

CUDA 内存不足：默认的 batch_size 已经是 1。请在数据增强配置中降低图像分辨率，或使用更小的 Swin 变体。

全景格式与实例格式不匹配：请确保提供的标注格式与 model.mode 设置匹配。

顶层 dataset.type 的部署架构错误：TAO Deploy 使用 dataset.val.type 和 dataset.test.type。请勿在 Mask2Former 部署规范的 顶层设置 dataset.type。

极小分辨率下的 ONNX 导出形状断言错误：如果 PyTorch ONNX 形状推断因 minus_one_pos != -1 而导致导出失败，请在重试部署验证前恢复模板导出尺寸（960x544）。必要时，可使用较小的训练和评估图像尺寸进行快速冒烟测试，但除非已验证目标形状，否则不要在导出时使用这些极小尺寸。

量化检查点加载错误：较旧的 PyTorch 镜像可能会导致基于检查点的 mask2former quantize 失败，因为运行时量化脚本向 Mask2formerPlModule.load_from_checkpoint 传递了 experiment_spec，而不是必需的 cfg 参数。包含量化修复的镜像支持默认的 torchao 检查点流程。ONNX 量化仍需设置 backend: modelopt.onnx、mode: static_ptq、固定的 dataset.test.target_size，并使用包含 modelopt.onnx.quantization 的镜像。

Spec 参数 / 父模型推断

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

来自 TAO Core mask2former.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	当前作业结果目录
生成 TRT 引擎	encryption_key	key	加密密钥
生成 TRT 引擎	gen_trt_engine.onnx_file	parent_model	从父作业结果文件夹推断出的模型文件
生成 TRT 引擎	gen_trt_engine.trt_engine	create_engine_file	输出 TensorRT 引擎路径
生成 TRT 引擎	results_dir	output_dir	当前作业结果目录
推理	encryption_key	key	加密密钥
推理	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
推理	inference.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
推理	results_dir	output_dir	当前作业结果目录
量化	encryption_key	key	加密密钥
量化	quantize.model_path	parent_model	从父作业结果文件夹推断出的模型文件
量化	results_dir	output_dir	当前作业结果目录
训练	encryption_key	key	加密密钥
训练	model.backbone.pretrained_weights	{'link': 'https://github.com/SwinTransformer/storage/releases/download/v1.0.8/swin_tiny_patch4_window7_224_22k.pth', 'destination_path': '/ptm/mask2former/swin_tiny_patch4_window7_224_22k/swin_tiny_patch4_window7_224_22k.pth'}	{'link': 'https://github.com/SwinTransformer/storage/releases/download/v1.0.8/swin_tiny_patch4_window7_224_22k.pth', 'destination_path': '/ptm/mask2former/swin_tiny_patch4_window7_224_22k/swin_tiny_patch4_window7_224_22k.pth'}
训练	results_dir	output_dir	当前作业结果目录
训练	train.resume_training_checkpoint_path	resume_model	从当前作业结果文件夹推断出的模型文件

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

在 SDK 解析器之外选择 Mask2Former 检查点时，应精确匹配目标 epoch/step 产物，例如 model_epoch_000_step_00100.pth。只有在明确请求最新版本时，mask2former_model_latest.pth 符号链接才有效。

部署

• tao-deploy-mask2former