tao-train-segformer

SegFormer

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

SegFormer 用于语义分割。它采用轻量级的基于 Transformer 的架构，并通过层级式特征提取提高效率，适用于实时分割任务。

设置 model.backbone.pretrained_backbone_path 以指定骨干网络权重。

对于 TAO Deploy TensorRT 操作（gen_trt_engine、TensorRT evaluate 和 TensorRT inference），请先阅读 references/tao-deploy-segformer.md。部署规范模板位于此 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 覆盖值不会更改模型元数据。

支持的操作

已打包的 SegFormer PyT CLI 支持 train、evaluate、export、inference、quantize 和 default_specs。此模型 skill 提供 train、evaluate、export、inference 和 quantize；恢复/重新训练通过 train 配合 train.resume_training_checkpoint_path 执行。

上层 PyT CLI 不提供 gen_trt_engine。请使用 models/segformer/deploy 生成 TensorRT 引擎、执行 TensorRT 评估和 TensorRT 推理。

训练要求

• 数据集类型： 分割
• 格式： unet
• AutoML 训练指标： val_miou（最大化）。
• 独立评估指标： test_miou。AutoML 排名使用训练状态中的 val_miou；评估器的 test_miou 仅用于检查点验证。

每个操作的数据集要求

操作	Spec 键	来源	文件	是否为列表？
evaluate	dataset.segment.root_dir	eval_dataset	解压后的根目录，其中包含 images/<split> 和 masks/<split>	否
export	dataset.segment.root_dir	train_datasets	解压后的根目录，其中包含 images/<split> 和 masks/<split>	否
inference	dataset.segment.root_dir	inference_dataset	解压后的根目录，其中包含 images/<split> 和 masks/<split>	否
quantize	dataset.segment.root_dir	train_datasets	解压后的根目录，其中包含 images/<split> 和 masks/<split>	否
quantize	dataset.segment.quant_calibration_dataset.images_dir	calibration_dataset	解压后的图像目录	否
train	dataset.segment.root_dir	train_datasets	解压后的根目录，其中包含 images/<split> 和 masks/<split>	否

常见 Spec 覆盖项

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

SEG_TRAIN_ROOT = "/data/segformer/train"
SEG_EVAL_ROOT = "/data/segformer/eval"
SEG_INFER_ROOT = "/data/segformer/infer"
CAL_IMAGES = f"{SEG_TRAIN_ROOT}/images/train"

train（必需数据源）：

{
    "train.num_gpus": 1,
    "train.num_epochs": 10,
    "train.checkpoint_interval": 10,
    "train.validation_interval": 10,
    "dataset.segment.batch_size": 4,
    "dataset.segment.root_dir": SEG_TRAIN_ROOT,
}

evaluate（必需数据源）：

{
    "evaluate.batch_size": 4,
    "dataset.segment.root_dir": SEG_EVAL_ROOT,
    "evaluate.checkpoint": CHECKPOINT,
}

inference（必需数据源）：

{
    "dataset.segment.batch_size": 1,
    "dataset.segment.root_dir": SEG_INFER_ROOT,
    "inference.checkpoint": CHECKPOINT,
}

export（必需数据源）：

{
    "dataset.segment.root_dir": SEG_TRAIN_ROOT,
    "export.checkpoint": CHECKPOINT,
    "export.input_height": 256,
    "export.input_width": 256,
    "export.onnx_file": ONNX_FILE,
}

quantize（必需数据源）：

{
    "dataset.segment.root_dir": SEG_TRAIN_ROOT,
    "dataset.segment.quant_calibration_dataset.images_dir": CAL_IMAGES,
    "quantize.model_path": CHECKPOINT,
}

如果源数据集以单独的 images/*.tar.gz 和 masks/*.tar.gz 压缩包提供，请在启动前解压，以便 root_dir 包含 images/train、images/val、images/test、masks/train 和 masks/val 等目录。不要将 dataset.segment.root_dir 指向仍只包含 tar 包的临时解压目录。

评估数据集

可选。验证数据通常包含在 root_dir 结构中。

重要参数

• dataset.segment.num_classes：分割类别数。默认值为 2（二分类）。必须与掩码标注中的类别数一致。
• model.backbone.type：默认值为 fan_small_12_p4_hybrid。支持 FAN 变体、SegFormer MIT 变体及其他类型。
• dataset.segment.root_dir：分割数据集的根目录。
• dataset.segment.img_size：输入图像尺寸。默认值为 256。增大尺寸可获得更精细的分割效果，但会增加内存开销。
• train.optim.lr：学习率。默认值为 6e-5。
• model.freeze_backbone：训练期间是否冻结骨干网络。适用于使用有限数据进行微调。
• dataset.segment.batch_size：每 GPU 的批次大小。默认值为 8。
• dataset.segment.label_transform：不需要标签变换时，使用字符串 "None"。不要将其设为 JSON/YAML null；严格模式的 schema 合并会将此字段视为字符串枚举。
• dataset.segment.palette：对于灰度掩码，每个 RGB 条目使用一个整数，例如 rgb: [85]。保留数据集的实际标签 ID 和类别名称，除非用户明确要求转换，否则不要对其进行标准化。

多 GPU / 多节点

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

Spec 键	说明	默认值
train.num_gpus	GPU 数量	1
train.gpu_ids	GPU 设备索引	[0]
train.num_nodes	节点数量	1
train.sync_batchnorm	在 GPU 间同步 BN	可配置
train.use_distributed_sampler	使用分布式采样器	可配置

• 多 GPU 策略：ddp_find_unused_parameters_true
• 不支持 FSDP

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

硬件

最低需要 1 块 GPU，建议使用 2 块 GPU。每块 GPU 需要 16GB 或以上显存（V100 或 A100）。SegFormer 相对轻量。默认 img_size=256，对内存较为友好。增大 img_size 可提高分辨率，但会增加内存开销并降低速度。

错误模式

CUDA 内存不足：减小 batch_size 或 img_size。SegFormer 的内存占用随图像尺寸呈平方级增长。

num_classes 不匹配：确保 dataset.segment.num_classes 与掩码标注中的实际类别数一致。

分割训练不支持 TensorBoard：保持 train.tensorboard.enabled: false。SegFormer 训练入口会断言分割任务不支持 TensorBoard 可视化，因此不要仅为提取 AutoML 指标而启用 TensorBoard；请改用日志解析或训练后的评估器。

AutoML 指标提取：SegFormer 训练状态文件会报告 val_miou，以及 val_loss、val_acc 和其他验证 KPI。默认的 AutoML 训练启动必须以 val_miou 为优化目标，并将 direction 设为 maximize；默认模型调用不得优化 val_loss。

对于 AutoML 或较长的分割搜索，请优先从 results_dir/train/status.json 中读取 val_miou。如果包装器报告终止失败，但结构化状态文件已达到配置的训练预算，且包含有限的 val_miou，则应报告恢复的指标，并注明包装器失败，而不是丢弃该测量结果。

对于高分辨率自定义分割目标，请将数据集路径作为每次运行的输入。不要将客户或用户专用的根路径添加到此可复用技能中。当用户要求固定的完整预算搜索时，请记住，括号算法（asha、bohb、dehb、hyperband、hyperband_es、pbt）可能会有意降低某些建议的 train.num_epochs；如果每条建议都必须运行完整的轮数，请使用 Bayesian/BFBO，或锁定预算。

检查点交接：对于评估、导出、推理、量化或恢复操作，请使用检查点解析器，从最佳 AutoML 子作业的 results_dir/train/ 文件夹中选择适合该操作的 model_epoch_*.pth 检查点，例如 model_epoch_000_step_00010.pth。SegFormer 也可能写入 segformer_model_latest.pth，但只有调用方明确请求最新检查点时才应使用该文件。为后续操作保留 dataset.segment.num_classes、dataset.segment.img_size 和 dataset.segment.root_dir 覆盖项。

恢复／重新训练检查点：恢复训练使用 train.resume_training_checkpoint_path。 请传入从先前训练输出中解析出的确切检查点，不要猜测 model.pth 路径。恢复后运行 1 个轮次时，应在新的结果目录中生成下一个检查点，例如 model_epoch_001_step_00020.pth。

导出／TensorRT 形状对齐：除非已在其他分辨率下验证训练模型和部署规格，否则请保持 export.input_height 和 export.input_width 与 dataset.segment.img_size 一致。打包后的全新安装流程已在 256x256 下验证，与默认的 SegFormer 数据集和部署模板相匹配。

PyT CLI 拒绝父级 segformer gen_trt_engine：在经过验证的 7.0.0 PyT 容器中，segformer gen_trt_engine 不是有效的父模型子任务。请使用 SegFormer 部署工作流（references/tao-deploy-segformer.md）生成 TensorRT 引擎、执行 TensorRT 评估和 TensorRT 推理。

Spec 参数 / 父模型推断

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

TAO Core segformer.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	当前作业的结果目录
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	model.backbone.pretrained_backbone_path	ptm_if_no_resume_model	不存在恢复检查点时使用的 PTM
train	results_dir	output_dir	当前作业的结果目录
train	train.pretrained_model_path	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-segformer