SuperAPI AI 工具集

tao-train-segformer

通用与其它

查看项目

更新于 10-16 08:11

技能介绍

用途:在 TAO 技能库中完成 SegFormer 语义分割模型的训练与后续流程,覆盖 train、evaluate、inference、export、quantize 以及 gen_trt_engine 等 TensorRT 部署动作。

主要能力:SegFormer 采用轻量级 Transformer 分层特征提取,适合实时分割任务。数据集类型为 segmentation,格式 unet,监控指标 val_miou。关键参数包括 dataset.segment.num_classes(默认 2,需与掩码类别数一致)、model.backbone.type(默认 fan_small_12_p4_hybrid,支持 FAN 与 SegFormer MIT 等变体)、dataset.segment.img_size(默认 256)、train.optim.lr(默认 6e-5)、model.freeze_backbone、dataset.segment.batch_size(默认 8)。骨干权重通过 model.backbone.pretrained_backbone_path 设置。

使用方式:每个动作都必须构造数据源路径并写入 spec_overrides,例如 train 使用 dataset.segment.root_dir 指向训练数据,evaluate 与 inference 指向评估数据,quantize 还需设置 dataset.segment.quant_calibration_dataset.images_dir。多 GPU 由 Lightning 管理,默认 train.num_gpus=1,策略为 ddp_find_unused_parameters_true,不支持 fsdp;多节点环境变量由编排器设置。训练阶段若 automl_policy 为 auto 且 automl_enabled 为 true 且打包了 train schema 与模板,默认经 tao-skill-bank:tao-run-automl 路由;用户要求关闭 AutoML 时按 off 处理,仅走直接训练。

适用场景与限制:适合需要实时语义分割的训练与部署流水线。最低 1 张 GPU,推荐 2 张,每卡 16GB 以上显存(V100 或 A100)。显存不足时降低 batch_size 或 img_size,SegFormer 显存随图像尺寸平方增长。TensorRT 相关动作需先阅读 references/tao-deploy-segformer.md。运行时不存在 ~/tao-core,schema 与模板由维护者在打包前生成。若缺少 train schema 或模板,需说明 AutoML 已启用但当前不可运行。其他未在资料中说明的收费、平台与效果信息暂未提供。