⚡

TAOSegFormer语义分割训练Skill tao-train-segformer

该技能用于在NVIDIA TAO工具套件中训练和部署SegFormer轻量级Transformer语义分割模型，支持从数据准备、分布式训练（含AutoML自动调参）到评估、导出、量化和TensorRT推理的完整流程。关键词：SegFormer；语义分割；TAO训练；轻量级Transformer；实时分割；模型训练；模型评估；模型导出；INT8量化；TensorRT。

视觉模型训练 0 次安装 1 次浏览 更新于 9/6/2026
名称 tao-train-segformer
描述 SegFormer用于语义分割。基于轻量级Transformer架构，具有分层特征提取能力，适用于实时分割任务。在需要对TAO SegFormer模型进行训练、评估、导出、量化或推理时使用。触发短语包括“train SegFormer”、“semantic segmentation”、“lightweight transformer segmenter”、“real-time semantic segmentation”。
开源协议 Apache-2.0 compatibility: 需要docker + nvidia-container-toolkit。 metadata:
版本 “0.1.0”
作者 NVIDIA Corporation allowed-tools: Read Bash tags: - segmentation

SegFormer

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

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。

Dataclass Schemas

生成的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；维护人员在打包技能库之前会重新生成模式和模板。

Train Action Policy

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

Supported Actions

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

父级PyT CLI不公开gen_trt_engine。使用models/segformer/deploy进行TensorRT引擎生成、TensorRT评估和TensorRT推理。

Training Requirements

  • 数据集类型: segmentation
  • 格式: unet
  • 监控指标: val_miou，最大化

Per-Action Dataset Requirements

操作 规范键 来源 文件 列表?
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> 否

Typical Spec Overrides

数据源覆盖对所有操作是必填的——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指向仍只包含压缩包的存档暂存文件夹。

Eval Dataset

可选。验证数据通常是root_dir结构的一部分。

Important Parameters

  • 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；严格的模式合并会将该字段视为字符串枚举。
  • dataset.segment.palette: 对于灰度掩码，每个RGB条目使用一个整数，例如rgb: [85]。请保留数据集的真实标签ID和类别名称，除非用户明确要求转换。

Multi-GPU / Multi-Node

启动方式： Lightning管理（单个python进程，Lightning产生worker）。

规范键 描述 默认
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。

Hardware

最少1个GPU，推荐2个GPU。每个GPU 16GB+（V100或A100）VRAM。SegFormer相对轻量。默认img_size=256对内存友好。增加img_size可提高分辨率，但会增加内存和速度开销。

Error Patterns

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，方向为maximize；不要为默认模型调用优化val_loss。

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

对于高分辨率自定义分割目标，保持数据集路径作为每次运行的输入。不要向这个可重用技能添加客户/用户特定的根目录。当用户请求固定的完整预算搜索时，记住括号算法（asha、bohb、dehb、hyperband、hyperband_es、pbt）可能会故意降低某些建议的train.num_epochs；如果每个建议都必须运行完整的epoch数，请使用贝叶斯/BFBO或锁定预算。

检查点交接: 对于evaluate/export/inference/quantize/resume，使用检查点解析器在最佳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路径。恢复的一次epoch运行应该在新结果目录中产生下一个检查点，例如model_epoch_001_step_00020.pth。

导出/TensorRT形状对齐: 保持export.input_height和export.input_width与dataset.segment.img_size一致，除非训练好的模型和部署规范已在另一个分辨率下验证。打包的全新安装路径在256x256下验证，与默认SegFormer数据集和部署模板匹配。

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

Spec Param / Parent Model Inference

模型特定的推断映射属于此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 当前作业结果目录
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 当不存在恢复检查点时的预训练模型
train results_dir output_dir 当前作业结果目录
train train.pretrained_model_path ptm_if_no_resume_model 当不存在恢复检查点时的预训练模型
train train.resume_training_checkpoint_path resume_model 从当前作业结果文件夹推断的模型文件

对于parent_model或parent_model_folder，请将上游train/export/AutoML子作业ID作为parent_job_id传递。SDK列出父结果文件夹，过滤检查点工件，并返回所选模型文件或文件夹。不要将这些映射添加回config.json，也不要修补生成的运行器脚本以猜测检查点路径。

Deployment