tao-train-segformer

SegFormer

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

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

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

训练要求

• 数据集类型： 分割
• 格式： unet
• 监控指标： val_miou

各操作的数据集要求

操作	规范键	来源	文件	列表？
evaluate	dataset.segment.root_dir	eval_dataset		否
export	dataset.segment.root_dir	train_datasets		否
inference	dataset.segment.root_dir	eval_dataset		否
quantize	dataset.segment.root_dir	train_datasets		否
quantize	dataset.segment.quant_calibration_dataset.images_dir	train_datasets		否
train	dataset.segment.root_dir	train_datasets		否

典型规范覆盖设置

每项操作都必须覆盖数据源——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,
    "dataset.segment.batch_size": 4,
    "dataset.segment.root_dir": f"{S3_TRAIN}",
}

evaluate（必需数据源）：

{
    "evaluate.batch_size": 4,
    "dataset.segment.root_dir": f"{S3_EVAL}",
}

gen_trt_engine：

{
    "gen_trt_engine.tensorrt.data_type": "fp16",
}

inference（必需数据源）：

{
    "dataset.segment.batch_size": 1,
    "dataset.segment.root_dir": f"{S3_EVAL}",
}

export（必需数据源）：

{
    "dataset.segment.root_dir": f"{S3_TRAIN}",
}

quantize（必需数据源）：

{
    "dataset.segment.root_dir": f"{S3_TRAIN}",
    "dataset.segment.quant_calibration_dataset.images_dir": f"{S3_TRAIN}",
}

评估数据集

可选。验证数据通常包含在 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。

多 GPU / 多节点

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

Spec Key	Description	Default
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 与掩码标注中的实际类别数一致。

规范参数 / 父模型推断

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

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

操作	规范字段	推断函数	含义
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	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 — 使用 TAO Deploy 生成 TensorRT 引擎、进行 TensorRT 评估和 TensorRT 推理的 SegFormer 部署工作流。