先说我最近做的这个事。手头攒了一批高分遥感影像要做建筑轮廓自动分割最早用U-Net和DeepLabV3边缘细节始终差一口气。后来在MMSegmentation里试了SegformermIoU直接涨了六七个点而且训练配置比想象中简单只是中间踩了不少版本、数据、显存相关的坑。这篇文章就把完整流程写清楚怎么从零开始用Segformer训练自己的数据包括环境搭建、数据标注与转换、配置文件逐段解读、训练监控、评估可视化以及最容易绊倒新手的几个问题。内容偏实操尽量每一步都能照着做。1. Segformer的四个关键设计以及它凭什么适合自定义数据集1.1 分层Transformer与重叠Patch MergingSegformer不是像ViT那样把整张图切成固定大小的patch一次性塞进Transformer而是走了一条和Swin类似的分层路线输入图片先切成4x4的patch然后通过多层Encoder逐级合并输出分辨率从1/4一路降到1/32形成特征金字塔。这个设计对分割任务至关重要。分割既需要高分辨率特征来保留边缘细节也需要低分辨率特征来提供全局语义。Segformer天然输出多尺度特征配合后面的轻量解码器正好覆盖这两种需求。Patch Merging在做合并时有一个细节不是简单地把2x2的patch拼接而是用重叠的卷积核来合并相邻patch之间有重叠区域。这样做的好处是能减少“方块感”避免最后的分割结果出现明显的格子纹理。我自己在遥感建筑分割里对这一点体会很深——一旦patch边界处理不好楼房屋顶边缘会出现很规则的小锯齿。1.2 Efficient Self-Attention与MLP DecoderTransformer最大的痛点就是自注意力机制的复杂度是O(N²)在高分辨率分割任务上显存很容易爆炸。Segformer的解法是在计算注意力之前先对Q和V做序列降采样实际效果等效于把复杂度降到了接近线性。这也是为什么Segformer能比较轻松地吃下512x512甚至更大尺寸的训练裁剪块。解码器部分就更直接了把backbone输出的四层特征全部上采样到1/4分辨率拼接起来过一个MLP就得到分割结果。没有ASPP、没有FPN简单得让第一次用的人有点怀疑。但仔细想想就明白Transformer层已经把全局上下文编码在特征里了不再需要一个复杂解码器去补感受野。这个设计还有一个额外好处解码器参数很少训练时不容易过拟合在小数据集上同样稳。1.3 什么时候应该换别的模型Segformer确实好但不是无脑选。如果你的任务只有几十张训练图、两三个类别而且对实时性要求极高那轻量级模型比如STDC或者MobileNetV3DeepLabV3会更合适。反过来如果你的数据量有几百上千张类别复杂尤其是遥感、医学影像这类要求精细边缘的任务Segformer的预训练权重迁移优势就非常明显。另外要注意Segformer官方预训练权重是在ImageNet上做的输入的mean/std也是ImageNet统计出来的那套。如果你用的是16bit卫星影像或者多波段数据就要额外做一些标准化处理不能直接把官方数据预处理搬过来用。2. 环境搭建conda、PyTorch、mmcv三件套的版本匹配实战2.1 我实测的版本组合Segformer可以用官方源码跑也可以直接在MMSegmentation里跑。我推荐后者因为数据加载、日志、评估、分布式训练都帮你封装好了省下大量造轮子的时间。我的环境组合如下实测稳定组件版本Python3.9PyTorch2.0.1torchvision0.15.2CUDA11.8驱动支持即可mmcv2.0.1mmsegmentation1.2.2这组版本不是唯一的答案但我在多台机器上装过这套最省心。如果你之前已经装了其他版本也没必要全部卸载重来可以参考官方表格找对应关系。2.2 安装命令与坑位说明先创建虚拟环境避免和已有项目打架conda create -n seg python3.9 -y conda activate seg然后装PyTorch。这里的关键是CUDA版本必须和你的显卡驱动以及后续的mmcv匹配pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118装完先别急着装mmcv我用的是openmim来装它会自动探测当前环境的PyTorch和CUDA版本拉取对应的预编译包pip install -U openmim mim install mmcv2.0.1很多人卡在mmcv这一关是因为直接用pip install mmcv结果它开始现场编译然后因为GCC版本、CUDA_HOME路径之类的编译问题失败。用mim装预编译包基本能绕开这个问题。接着从GitHub拉取MMSegmentation并进入仓库目录git clone https://github.com/open-mmlab/mmsegmentation.git cd mmsegmentation pip install -e .2.3 环境验证装完之后别急着开始训练先快速验证一下整套环境能不能跑通。去MMSegmentation的model zoo下载Segformer B0在Cityscapes上的预训练权重然后跑一次测试python tools/test.py configs/segformer/segformer_mit-b0_8xb1-160k_cityscapes-1024x1024.py \ /path/to/segformer_mit-b0_512x512_160k_ade20k.pth \ --eval mIoU能正常输出mIoU数据就说明环境没问题。这一步看起来多花了几分钟实际能帮你省下后面排查“为啥一训练就报错”的大量时间。3. 数据准备从标注到MMSegmentation标准数据集的完整转换3.1 标注工具怎么选我建议中小数据集用LabelMe因为安装简单、JSON格式直观、官方转换工具多。团队协作场景用CVAT支持在线标注、多人协作导出的格式也比较标准。如果你预算充足X-AnyLabeling这类半自动标注工具能先跑一个Segformer预训练模型做预标注人工只负责修正效率提升非常明显。标注时有一个容易忽视的原则背景不用单独画polygonmask像素值为0的默认就是背景。类别标签尽量建少一点宁可后期合并也别一开始就建十几类然后标得乱七八糟。3.2 从LabelMe JSON到标准mask的转换LabelMe保存的JSON包含每个shape的标签和多边形点坐标。转换逻辑很简单遍历所有shape在空白mask上用对应类别索引值填充多边形区域。下面这个脚本我一直在用注释写得很清楚import json import glob import os import numpy as np from PIL import Image, ImageDraw # 类别名与索引的映射背景固定为0 class_names [building, road] label_map {name: i 1 for i, name in enumerate(class_names)} os.makedirs(masks, exist_okTrue) for json_file in glob.glob(labels/*.json): with open(json_file, r, encodingutf-8) as f: data json.load(f) h data[imageHeight] w data[imageWidth] mask Image.new(L, (w, h), 0) for shape in data[shapes]: cls shape[label] if cls not in label_map: continue pts [tuple(point) for point in shape[points]] ImageDraw.Draw(mask).polygon(pts, outlinelabel_map[cls], filllabel_map[cls]) base os.path.splitext(os.path.basename(json_file))[0] mask.save(os.path.join(masks, base .png))这里的核心要点是保存出来的mask必须是单通道灰度PNG而不是三通道RGB彩色图。训练时模型拿的是像素值索引不是颜色。如果你保存的时候不小心转成了RGB后面训练时会报shape mismatch或者loss莫名其妙降不下去。3.3 目录组织与自定义Dataset类MMSegmentation支持直接以文件夹形式读取数据不需要转成lmdb或者TFRecord。目录按照下面的结构组织data/my_dataset/ ├── img_dir/ │ ├── train/ │ │ ├── img_001.png │ │ └── ... │ └── val/ │ └── ... └── ann_dir/ ├── train/ │ ├── img_001.png │ └── ... └── val/ └── ...然后写一个自定义Dataset类注册进MMSegmentationimport os import numpy as np from mmseg.registry import DATASETS from mmseg.datasets.basesegdataset import BaseSegDataset DATASETS.register_module() class MyDataset(BaseSegDataset): METAINFO dict( classes(background, building, road), palette[[0, 0, 0], [0, 255, 0], [255, 0, 0]] ) def __init__(self, **kwargs): super().__init__( img_suffix.png, seg_map_suffix.png, **kwargs )这个文件保存为my_dataset.py。METAINFO里的classes顺序必须和转换脚本里label_map的索引对应palette只是给可视化用的不影响训练。训练和验证的图像名要一致程序通过文件名去匹配img和ann这对对应关系。3.4 遥感大图的切块与预处理如果你做的是遥感分割原始影像动辄上万像素不能直接整图进网络。我的做法是切成1024x1024的块相邻块之间保留128像素重叠。重叠区域在预测拼接时直接丢弃能明显减少边缘伪影。一个经验是先切块再标注而不是先标注再切块。但实际标注时目标对象往往跨块标注过程比较痛苦。如果你是先标注大图再切块切块脚本要注意mask和被切图像保持相同坐标系别把标注错位了。切完之后随机抽几张把mask叠加到图像上人工看一遍确认没有偏移再进入训练。4. 训练配置文件逐段拆解改哪里、为什么改、改了有什么后果4.1 配置的核心心智模型MMSegmentation把训练需要的所有东西都拆成了dict数据怎么加载、模型长什么样、优化器用什么、训练多少轮。你需要改的只有几个地方其他官方默认配置就别动了。下面是去掉注释也能直接用的一份独立config基于B0模型二分类任务支持512x512训练norm_cfg dict(typeBN, requires_gradTrue) data_preprocessor dict( typeSegDataPreProcessor, mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375], bgr_to_rgbTrue, pad_val0, seg_pad_val255) model dict( typeEncoderDecoder, data_preprocessordata_preprocessor, backbonedict( typeMixVisionTransformer, in_channels3, embed_dims[64, 128, 320, 512], num_heads[1, 2, 5, 8], num_layers[2, 2, 2, 2], out_indices(0, 1, 2, 3), drop_rate0.0, drop_path_rate0.1, pretrained/path/to/mit_b0.pth), decode_headdict( typeSegformerHead, in_channels[64, 128, 320, 512], in_index[0, 1, 2, 3], channels256, num_classes2, loss_decodedict( typeCrossEntropyLoss, use_sigmoidFalse, loss_weight1.0)), train_cfgdict(), test_cfgdict(modewhole)) dataset_type MyDataset data_root data/my_dataset/ train_pipeline [ dict(typeLoadImageFromFile), dict(typeLoadAnnotations), dict(typeRandomResize, scale(512, 1024), ratio_range(0.5, 2.0), keep_ratioTrue), dict(typeRandomCrop, crop_size(512, 512), cat_max_ratio0.75), dict(typeRandomFlip, prob0.5), dict(typePackSegInputs) ] test_pipeline [ dict(typeLoadImageFromFile), dict(typeLoadAnnotations), dict(typePackSegInputs) ] train_dataloader dict( batch_size8, num_workers6, persistent_workersTrue, samplerdict(typeDefaultSampler, shuffleTrue), datasetdict( typedataset_type, data_rootdata_root, data_prefixdict( img_pathimg_dir/train, seg_map_pathann_dir/train), pipelinetrain_pipeline)) val_dataloader dict( batch_size1, num_workers4, persistent_workersTrue, samplerdict(typeDefaultSampler, shuffleFalse), datasetdict( typedataset_type, data_rootdata_root, data_prefixdict( img_pathimg_dir/val, seg_map_pathann_dir/val), pipelinetest_pipeline)) test_dataloader val_dataloader val_evaluator dict(typeIoUMetric, iou_metrics[mIoU]) test_evaluator val_evaluator optim_wrapper dict( typeOptimWrapper, optimizerdict( typeAdamW, lr0.00006, betas(0.9, 0.999), weight_decay0.01), paramwise_cfgdict(custom_keys{ pos_block: dict(decay_mult0.), norm: dict(decay_mult0.) })) param_scheduler [ dict(typeLinearLR, start_factor1e-6, by_epochFalse, begin0, end1500), dict(typePolyLR, eta_min0.0, power1.0, by_epochFalse, begin1500, end160000) ] train_cfg dict( typeIterBasedTrainLoop, max_iters160000, val_interval4000) val_cfg dict(typeValLoop) test_cfg dict(typeTestLoop)4.2 每个关键字段为什么要这么设backbone部分如果用的是B0embed_dims是[64, 128, 320, 512]num_layers是[2, 2, 2, 2]。这里最重要的是pretrained字段指向你下载的mit_b0.pth。这个预训练权重对最终效果的影响非常大尤其是数据量不够大的时候别省这一步。decode_head的num_classes必须改成你的类别数。这个数字是背景加所有前景类别的总数。我见过有人二分类任务写了2结果background算一类的意思应该是2没问题但也有人三分类任务写成2导致最后训练不报错、评估全错。一定要数清楚如果标签有background、building、road三类num_classes就是3。optimizer用AdamW学习率6e-5。这个数值是用官方预训练权重做微调时的安全值。如果从零训练可以放到1e-4甚至更高但分割任务几乎没人从零训Transformer。LinearLR预热1500个iter是为了让模型在初期别因为学习率过大把预训练权重冲坏。4.3 RandomCrop的cat_max_ratio是什么这个参数在遥感分割里很有用如果一张子图里某个类别的像素占比超过0.75就重新裁剪。遥感影像经常出现建筑密集区模型会倾向于把所有东西都预测成建筑cat_max_ratio起到一个强行平衡类别的作用。如果你的数据是道路、屋顶这类长条形目标crop_size可以再调大一点到640或768但显存会跟着涨后面我会讲对应的显存优化方案。5. 训练启动与过程监控单卡多卡命令、显存应对和日志解读5.1 训练命令先确保你进入了mmsegmentation仓库目录并且把config文件和my_dataset.py放在正确位置。单卡训练命令python tools/train.py my_segformer_config.py --work-dir work_dirs/segformer_b0_my_data多卡训练用官方脚本这里的2是卡数CUDA_VISIBLE_DEVICES0,1 bash tools/dist_train.sh my_segformer_config.py 2 \ --work-dir work_dirs/segformer_b0_my_data--work-dir会保存训练日志、checkpoint和可视化文件。建议每次实验都换一个独立目录避免不同实验的权重互相覆盖。5.2 训练日志的关键信息在哪里看训练过程中终端会实时打印loss、学习率、显存占用等信息。这里有个容易误读的点Segformer打印的loss是decode_head的CrossEntropyLoss它下降得很快前几百个iter就能降到0.5以下但这不代表模型已经收敛。更靠谱的判断方式是看验证集上的mIoU。我在config里设置了val_interval4000也就是每4000个iter跑一次验证。你在日志里搜mIoU关键字能看到类似这样的输出2024-05-20 15:32:10 - INFO - Iter [4000/160000] mIoU: 0.6120 mAcc: 0.7345 aAcc: 0.9021mIoU曲线进入平台期之后继续训练收益就很小了。我习惯把mIoU和loss的曲线画在一起看如果loss还在降但mIoU不涨多半是过拟合或者验证集数据分布有偏。5.3 显存不够时按什么顺序调整显存不足是训练Segformer最常遇到的问题我建议按下面的优先级来调优先级方法配置改动注意事项1减小crop_size512改为384感受野变小密集小目标效果可能下降2减小batch_size8改为4收敛可能变慢配合梯度累积3梯度累积每4步累积一次等效batch_size不变省显存4混合精度训练fp16开关显存几乎减半部分数据集精度略微下降MMSegmentation开启梯度累积和混合精度需要在optim_wrapper里做配置optim_wrapper dict( typeOptimWrapper, optimizerdict(typeAdamW, lr0.00006, weight_decay0.01), accumulative_counts4, mixed_precisiondict(typeAmpOptimWrapper) )注意累积步数增大的同时最好让val_interval也跟着增大否则验证频繁会增加额外显存压力。5.4 断点续训与load_from的区别训练中断后继续训练用python tools/train.py my_segformer_config.py --work-dir work_dirs/segformer_b0_my_data --resume auto这会自动读取work_dir里最新的checkpoint接着上次的iter继续跑。另外还有一个--cfg-options load_fromxxx.pth的用法它只是把权重加载进模型但不恢复优化器状态更适合用来做“换数据继续微调”和断点续训语义完全不同别搞混。6. 评估、可视化推理与模型导出6.1 用test脚本产出完整评估指标训练结束后先用官方test脚本做一次全量评估python tools/test.py my_segformer_config.py \ work_dirs/segformer_b0_my_data/latest.pth \ --out results.pkl --eval mIoU输出的指标很重要但很多新人只会看一个mIoU。我建议同时关注这几个指标含义注意事项aAcc全图像素准确率类别极度不均衡时会虚高mAcc各类别像素准确率均值能暴露占比较小类别的召回问题mIoU所有类别IoU均值综合效果训练时主要看这个每类别IoU单独输出重点看小目标和难分区域如果mIoU不低但建筑类别的IoU特别低那说明模型对建筑边缘的定位还是不准优先检查标注质量和数据增强。6.2 可视化推理脚本评估之后要做可视化肉眼检查分割结果。用MMSegmentation的高级API写一个批量预测脚本import os import cv2 import numpy as np from mmseg.apis import init_segmentor, inference_segmentor config_file my_segformer_config.py checkpoint_file work_dirs/segformer_b0_my_data/latest.pth model init_segmentor(config_file, checkpoint_file, devicecuda:0) os.makedirs(vis_results, exist_okTrue) palette np.array([[0, 0, 0], [0, 255, 0], [255, 0, 0]], dtypenp.uint8) for img_path in os.listdir(data/my_dataset/img_dir/val): img os.path.join(data/my_dataset/img_dir/val, img_path) result inference_segmentor(model, img) # result[0]的形状是(H, W)每个位置是类别索引 color_mask palette[result[0].astype(np.uint8)] # 叠加显示 original cv2.imread(img) overlay cv2.addWeighted(original, 0.5, color_mask, 0.5, 0) cv2.imwrite(os.path.join(vis_results, img_path), overlay)建议把每张结果图都仔细看一遍重点关注边缘是否贴合、小目标是否漏检、同类目标是否粘连。这些细节比mIoU数字更能反映模型能不能真正落地。6.3 导出ONNX的注意事项模型需要部署时再考虑导出。我的经验是别尝试把mmseg的整套推理流程导出直接导出backbone加decode_head的forward也就是输入一张归一化后的图像张量输出logits把前后处理留在部署框架里做。import torch model.eval() dummy torch.randn(1, 3, 512, 512).cuda() torch.onnx.export( model, dummy, segformer_b0.onnx, input_names[input], output_names[output], dynamic_axes{ input: {0: batch, 2: height, 3: width}, output: {0: batch, 2: height, 3: width} }, opset_version11 )导出后一定要用onnxruntime或者trtexec做一次输入输出对比检查最大误差。如果误差超过1e-3级别优先检查导出时是否混入了训练阶段的dropout或sync_bn。如果你要部署到TensorRT我更建议直接去看MMDeploy的官方教程它把预处理、归一化、后处理全部打包好了比自己瞎折腾稳定得多。7. 训练Segformer最容易踩的六个坑原因与排查链路7.1 mask不是单通道索引图导致loss不降这个坑的症状很典型训练loss一路降到0.3就不再动了验证mIoU始终在0.1以下预测结果基本是全黑或全图同一个类别。排查链路用PIL打开一张训练mask检查mode。如果输出是RGB说明你的mask保存成了三通道图。检查转换脚本是不是用了mask.convert(RGB)或者cv2.imread默认按三通道读取。在dataset类里打印一张样本的data_samples.gt_sem_seg.data.shape和torch.unique(...)确认值是离散的类别索引。修复方式转换脚本里一直使用Image.new(L, ...)保持单通道读取端确认seg_map的通道数为1。7.2 标签值不连续num_classes怎么算都不对有人标注时从1开始编号而不是从0开始比如背景是1、房子是2num_classes设成3。模型在计算loss时会把类别索引映射到维度导致某个维度永远训练不到最终这个类别的IoU是0。排查方法很简单把所有validation mask用numpy打开打印np.unique(mask)。如果最大值等于num_classes-1说明是正常的如果出现缺失的索引或者超出范围的值必须修正数据。7.3 加载预训练权重时报shape mismatch这是最常见的报错之一报错信息会显示size mismatch for decode_head.linear_fuse.weight。原因多半是decode_head的num_classes和你下载权重的类别数不一致或者backbone的pretrained路径写错了。处理方式先把num_classes改对然后backbone部分加载权重时使用strictFalse。在MMSegmentation里backbone的pretrained参数会自动处理这种非完全匹配情况一般不会卡在这。7.4 mmcv编译失败卡在GCC版本上如果你坚持不用mim而选择源码编译mmcv大概率会遇到这类问题。常见报错是gcc: error: unrecognized command line option或者找不到cuda_runtime.h。排查链路确认nvcc -V和python -c import torch; print(torch.version.cuda)一致PyTorch里的CUDA版本和系统CUDA Toolkit版本不一定要一致但编译时系统CUDA必须存在。设置export CUDA_HOME/usr/local/cuda很多GCC报错的根因其实是CUDA_HOME没指对。还不行的直接改用mim安装预编译包这是最省时间的路子。7.5 前几百个iter的loss不降反升这个现象我换数据集时遇到过几次原因几乎都是学习率太大。加载预训练权重后如果你把学习率设为1e-4甚至更高前几百个iter会把预训练特征冲乱loss会先上升再缓慢下降。我的建议是如果验证集mIoU在大概2000个iter之后还没恢复到初始水平停下来把lr降到3e-5重新启动。另外LinearLR预热不是可选项1500个iter的预热对这类型模型很重要。7.6 验证阶段OOM训练却没问题训练用crop_size裁剪后显存可控但验证时如果test_cfg设置的是modewhole模型会对整张图直接做推理。当你验证图是遥感大图时显存很容易直接爆掉。把config里的test_cfg改成滑窗推理test_cfg dict(modeslide, crop_size(512, 512), stride(384, 384))滑块之间保证有重叠重叠区域取置信度平均边缘伪影能明显减少。另外评估时batch_size保持1开persistent_workers可以稍微缓解反复新建进程的耗时。最后分享一个我自己的习惯每次训练结束我会把config文件、最后一次权重和评估pkl文件放到同一个work_dir下并手动在文件名里加上mIoU和训练iter数。后面调参时直接在标注了指标的目录之间对比比翻日志快得多。Segformer的训练流程本身并不复杂但数据质量、类别设置、学习率这些细节每一项都值得花时间认真对待。如果你想把它用到自己的数据上不妨先拿一个小数据集把全流程跑通再逐步扩展样本这是最快也最稳的路径。