tao-train-oneformer

OneFormer

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

OneFormer 用于通用图像分割。它通过使用任务条件查询的单一架构，统一全景分割、实例分割和语义分割。

设置 train.pretrained_backbone 和/或 train.pretrained_model。

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

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；维护者会在打包 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，但在生成架构前无法运行 AutoML。

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

训练要求

• 数据集类型： 分割
• 格式： coco_panoptic、coco
• AutoML 训练指标： mIoU，其中 direction=maximize
• 独立评估指标： test_mIoU，其中 direction=maximize

各操作的数据集要求

操作	Spec 键	来源	文件	是否为列表？
evaluate	dataset.train.images	train_datasets	images.tar.gz	否
evaluate	dataset.label_map	train_datasets	label_map.json	否
evaluate	dataset.train.annotations	train_datasets	annotations.json	否
evaluate	dataset.train.panoptic	train_datasets	images_panoptic.tar.gz	否
evaluate	dataset.val.images	eval_dataset	images.tar.gz	否
evaluate	dataset.val.annotations	eval_dataset	annotations.json	否
evaluate	dataset.val.panoptic	eval_dataset	images_panoptic.tar.gz	否
evaluate	dataset.test.images	eval_dataset	images.tar.gz	否
evaluate	dataset.test.annotations	eval_dataset	annotations.json	否
evaluate	dataset.test.panoptic	eval_dataset	images_panoptic.tar.gz	否
inference	dataset.train.images	train_datasets	images.tar.gz	否
inference	dataset.label_map	train_datasets	label_map.json	否
inference	dataset.train.annotations	train_datasets	annotations.json	否
inference	dataset.train.panoptic	train_datasets	images_panoptic.tar.gz	否
inference	dataset.val.images	eval_dataset	images.tar.gz	否
inference	dataset.val.annotations	eval_dataset	annotations.json	否
inference	dataset.val.panoptic	eval_dataset	images_panoptic.tar.gz	否
inference	dataset.test.images	inference_dataset	images.tar.gz	否
quantize	dataset.train.images	train_datasets	images.tar.gz	否
quantize	dataset.train.annotations	train_datasets	annotations.json	否
quantize	dataset.label_map	train_datasets	label_map.json	否
quantize	dataset.train.panoptic	train_datasets	images_panoptic.tar.gz	否
quantize	dataset.val.images	eval_dataset	images.tar.gz	否
quantize	dataset.val.annotations	eval_dataset	annotations.json	否
quantize	dataset.val.panoptic	eval_dataset	images_panoptic.tar.gz	否
quantize	dataset.test.images	eval_dataset	images.tar.gz	否
quantize	dataset.quant_calibration_dataset.images_dir	calibration_dataset	images.tar.gz	否
train	dataset.train.images	train_datasets	images.tar.gz	否
train	dataset.train.annotations	train_datasets	annotations.json	否
train	dataset.label_map	train_datasets	label_map.json	否
train	dataset.train.panoptic	train_datasets	images_panoptic.tar.gz	否
train	dataset.val.images	eval_dataset	images.tar.gz	否
train	dataset.val.annotations	eval_dataset	annotations.json	否
train	dataset.val.panoptic	eval_dataset	images_panoptic.tar.gz	否
train	dataset.test.images	eval_dataset	images.tar.gz	否

常见 Spec 覆盖项

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

S3_TRAIN = "s3://bucket/data/train"
S3_EVAL = "s3://bucket/data/eval"
S3_INFERENCE = "s3://bucket/data/inference"
S3_CALIBRATION = "s3://bucket/data/calibration"

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,
    "train.precision": "32",
    "dataset.train.images": f"{S3_TRAIN}/images.tar.gz",
    "dataset.train.annotations": f"{S3_TRAIN}/annotations.json",
    "dataset.label_map": f"{S3_TRAIN}/label_map.json",
    "dataset.train.panoptic": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.images": 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.images": f"{S3_TRAIN}/images.tar.gz",
    "dataset.label_map": f"{S3_TRAIN}/label_map.json",
    "dataset.train.annotations": f"{S3_TRAIN}/annotations.json",
    "dataset.train.panoptic": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.test.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.test.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
}

export：

{
    "export.checkpoint": "<selected train/AutoML checkpoint>",
    "model.sem_seg_head.num_classes": 133,
    "model.export": True,
    "export.onnx_file": "/results/oneformer_export_640.onnx",
}

inference（必需的数据源）：

{
    "inference.checkpoint": "<selected train/AutoML checkpoint>",
    "dataset.train.images": f"{S3_TRAIN}/images.tar.gz",
    "dataset.label_map": f"{S3_TRAIN}/label_map.json",
    "dataset.train.annotations": f"{S3_TRAIN}/annotations.json",
    "dataset.train.panoptic": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.images": f"{S3_INFERENCE}/images.tar.gz",
    "inference.images_dir": f"{S3_INFERENCE}/images.tar.gz",
}

quantize（必需的数据源）：

{
    "quantize.model_path": "<selected train/AutoML checkpoint>",
    "dataset.train.images": f"{S3_TRAIN}/images.tar.gz",
    "dataset.train.annotations": f"{S3_TRAIN}/annotations.json",
    "dataset.label_map": f"{S3_TRAIN}/label_map.json",
    "dataset.train.panoptic": f"{S3_TRAIN}/images_panoptic.tar.gz",
    "dataset.val.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.val.annotations": f"{S3_EVAL}/annotations.json",
    "dataset.val.panoptic": f"{S3_EVAL}/images_panoptic.tar.gz",
    "dataset.test.images": f"{S3_EVAL}/images.tar.gz",
    "dataset.quant_calibration_dataset.images_dir": f"{S3_CALIBRATION}/images.tar.gz",
}

检查点选择

OneFormer 训练会写入类似 model_epoch_000_step_00017.pth 的轮次-步数检查点，也可能会创建 oneformer_model_latest.pth 符号链接。对于依赖检查点的操作，请使用模型技能或 SDK 父模型解析器，并将确切选定的检查点路径传入 evaluate.checkpoint、inference.checkpoint、export.checkpoint、quantize.model_path 或 train.resume_training_checkpoint_path。除非用户明确要求使用最新检查点，否则不要仅凭名称选择 oneformer_model_latest.pth 符号链接。如果解析器报告了最佳检查点，则将该最佳检查点用于评估、导出或推理；如果用户指定了某个轮次或步数，则使用对应的轮次-步数检查点。

评估数据集

可选。验证数据与训练数据一同在数据集配置中设置。

重要参数

• model.sem_seg_head.num_classes：提供给预测头的分割类别索引数量。当 dataset.contiguous_id: True 通过标签映射重新映射原始类别 ID 时，COCO 全景数据的默认值为 133。除非标签映射和标注确实已缩减为该类别集，否则不要将此值缩减为全局工作流类别数。
• model.one_former.hidden_dim：本地冒烟测试时保持为 256，除非文本编码器维度也同步更改。仅降低 hidden_dim 会导致训练期间文本特征／上下文维度不匹配。
• model.backbone.name：默认值为 D2SwinTransformer（基于 Swin）。默认 embed_dim=192，depths=[2,2,18,2]。
• train.num_epochs：默认值为 50，明显高于大多数 TAO 模型。OneFormer 需要更多轮次才能收敛。
• train.optim.lr：学习率。默认值为 1e-5，低于 Mask2Former 的 2e-4。
• model.task_toggling：启用／禁用特定任务：semantic_on、instance_on、panoptic_on。
• export.task：导出任务模式。选项：semantic、instance、panoptic。默认值为 semantic。导出输入默认尺寸为 640x640。
• inference.mode：推理模式。选项：semantic、instance、panoptic。默认值为 semantic。image_size 的默认值为 [1024, 1024]。
• evaluate.iou_per_class：在评估中报告每个类别的 IoU。默认值为 True。

多 GPU／多节点

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

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

• 使用显式的 DDPStrategy，并设置 find_unused_parameters=True、gradient_as_bucket_view=True、process_group_backend="nccl"
• 始终启用 sync_batchnorm
• 不支持 fsdp；仅支持 DDP

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

硬件

至少需要 2 块 GPU，建议使用 4 块 GPU。每块 GPU 需要 24GB 以上显存（建议使用 A100）。OneFormer 和 Mask2Former 一样，对显存要求较高。默认 batch_size=1。要获得合理的训练速度，尤其是在训练 50 个轮次时，需要使用多 GPU。

错误模式

CUDA 内存不足：batch_size 已为 1。请降低图像分辨率或使用更小的 Swin 配置。

提取的 S3 tarball 路径高了一层：对于本地 Docker 运行， images.tar.gz 和 images_panoptic.tar.gz 解压后可能包含 images/ 和 images_panoptic/ 等包装目录。请将 dataset.*.images、 dataset.*.panoptic、inference.images_dir 和量化校准路径设置为实际包含图像或全景分割文件的文件夹，而不是包装目录。路径高了一层时，即使递归文件计数看起来正确，也会在读取第一张标注图像时触发 FileNotFoundError。

default_specs 缺少 results_dir：CLI default_specs 子任务会忽略 results_dir 的 -e 实验规格；请改用 Hydra 风格的覆盖项： oneformer default_specs results_dir=/path/to/default_specs。

Lightning 精度 fp32 无效：请在 train/AutoML/evaluate/inference 规格中使用 train.precision: "32"。当前 Lightning 技术栈不接受旧版 fp32 字符串。

PyTorch 2.6 检查点加载失败，影响后续操作：当前 OneFormer 检查点包含 OmegaConf 对象。对于由同一可信 TAO train/AutoML 工作流生成的检查点，请在后续 evaluate、inference、export、quantize 或 resume 作业的环境变量中设置 TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1，以便 Lightning 加载完整检查点。 请勿将此环境变量用于不可信的检查点。

matcher/class cost 中发生 CUDA device-side assert：如果训练在 oneformer/utils/matcher.py 中对 out_prob[:, tgt_ids] 进行索引时失败， 请将有效目标 ID 与 model.sem_seg_head.num_classes 进行比较。打包的 COCO panoptic 示例在 dataset.contiguous_id: True 重映射后包含 133 个紧凑类别，因此即使更宽泛的验证工作流传入了较小的通用 num_classes 值，也应使用 model.sem_seg_head.num_classes: 133。 只有在标签映射和标注已缩减为完全一致的连续类别集时，才使用较小的类别数。

推理返回 PASS，但没有预测结果：OneFormer 预测读取的是 inference.images_dir，而不是 dataset.test.images。每次运行推理时，都要声明并填充 inference.images_dir，指向图像文件夹或 tarball。dataset.test.images 仍可用于共享数据集上下文，但不会驱动 PyTorch predict 数据加载器。

导出输出路径预先创建为目录：请勿将 export.onnx_file 声明为文件输出。OneFormer 导出器会断言 ONNX 路径尚不存在，而本地运行器会预先创建已声明的输出路径。请在规格中明确将 export.onnx_file 设置为已挂载结果目录树下一个尚不存在的文件路径。烟雾验证时请保留默认的 640x640 导出形状；过小的导出形状可能触发 PyTorch ONNX 形状推断失败。

Quantize 无法从 AutoML 检查点找到训练标签映射： OneFormer Lightning 检查点会在保存的 hparams 中保留训练时的数据集绝对路径。从 AutoML 子检查点运行后续操作时，除了传入解析后的检查点路径，还要确保父级 AutoML 作业目录在操作容器内仍可通过其原始 /results/<job_id> 路径访问。否则，即使当前规格包含有效的 dataset.label_map，quantize 仍可能在加载检查点 hparams 时失败。

训练缓慢：单 GPU 上使用 batch_size=1 训练默认的 50 个 epoch 会很慢。请使用多 GPU 分布式训练。

Spec 参数／父模型推断

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

来自 TAO Core oneformer.config.json 的推断映射：

操作	Spec 字段	推断函数	含义
evaluate	encryption_key	key	加密密钥
evaluate	evaluate.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
evaluate	evaluate.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
evaluate	results_dir	output_dir	当前作业结果目录
export	encryption_key	key	加密密钥
export	export.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
export	export.onnx_file	create_onnx_file	输出 ONNX 路径
export	results_dir	output_dir	当前作业结果目录
gen_trt_engine	encryption_key	key	加密密钥
gen_trt_engine	gen_trt_engine.onnx_file	parent_model	从父作业结果文件夹推断出的模型文件
gen_trt_engine	gen_trt_engine.trt_engine	create_engine_file	输出 TensorRT 引擎路径
gen_trt_engine	results_dir	output_dir	当前作业结果目录
inference	encryption_key	key	加密密钥
inference	inference.checkpoint	parent_model	从父作业结果文件夹推断出的模型文件
inference	inference.trt_engine	parent_model	从父作业结果文件夹推断出的模型文件
inference	results_dir	output_dir	当前作业结果目录
quantize	encryption_key	key	加密密钥
quantize	quantize.model_path	parent_model	从父作业结果文件夹推断出的模型文件
quantize	results_dir	output_dir	当前作业结果目录
train	encryption_key	key	加密密钥
train	results_dir	output_dir	当前作业结果目录
train	train.pretrained_backbone	{'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'}
train	train.pretrained_model	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
train	train.resume_training_checkpoint_path	resume_model	从当前作业结果文件夹推断出的模型文件

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

部署

• tao-deploy-oneformer