在语义分割这个方向上我这两年用得最多的工具就是 Segformer。从最开始被它在 Cityscapes 上的精度吸引到后续在遥感、工业质检、自动驾驶辅助数据上反复折腾说实话踩过的坑不比跑通的流程少。今天这篇就不讲官网上那些复制粘贴的 example 了直接按我自己的实操路径把怎么用 Segformer 训练自己的数据完整拆一遍尤其是数据格式、训练参数、环境坑这三块基本面向想把自己标注好的数据跑出模型的同学。先简单交代一下它是什么、能干什么。Segformer 是那种 transformer 做编码器 轻量解码器的语义分割模型和 DeepLabv3 那一类纯卷积方案的最大区别在于它对全局上下文更敏感遇到细长目标、边缘模糊目标、多尺度目标混合的场景会稳一些。我主要用它跑遥感影像的建筑物提取和道路分割效果比之前用 deeplabv3 时明显省心不用死磕空洞卷积的 rate 组合。如果你是刚接触语义分割刚标注好一批数据想找个“能直接用、别太玄学”的模型这篇文章很适合你。就算你用的是 mmsegmentation 或 MMSeg很多内容也能直接对接上因为大部分训练逻辑是相通的。1. 内容整体设计与思路拆解1.1 为什么选 Segformer 而不是其他模型很多同学一上来就纠结“我该用 U-Net 还是 DeepLabv3 还是 Segformer”其实这个选择没那么复杂核心就看三件事数据规模、目标形态、部署环境。先说数据规模。Segformer 在 ImageNet 上预训练过编码器是 MiT 系列比如 B0、B1、B2 这些。B0 参数量很小只有 3.7M 左右用一小块显卡就能跑但小模型的感受野依然很大因为 transformer 的 self-attention 天然能建模长距离依赖。我拿标注好的 500 张遥感图做过对比同样迭代次数下Segformer-B0 的 mIoU 比 U-Net 高 4~6 个点左右而且不需要像 deeplabv3 那样手工设计不同的空洞卷积采样率省掉了最痛苦的调参过程。再说目标形态。如果你的目标物体是长条形道路、河流、碎块状建筑物、车、或者尺度差异很大大块农田小屋檐Segformer 的注意力机制能比较好地缓解尺度不匹配的问题。它的解码器不是简单上采样而是融合了不同层级的特征并做像素级加权这样高分辨率细节和低分辨率语义能同时被利用边缘不会糊成一团。最后是部署环境。Segformer-B0/B1 在 512x512 输入下单张 GTX 1660 就能训练推理速度也不慢。如果是 B3/B4就需要 24G 以上显存或者用混合精度。所以整体设计思路是先用 B0 跑通流程拿到 baseline再根据任务难度换更大的编码器。1.2 整体流程怎么规划我自己的固定流程是这五步后面每个环节都会展开数据准备统一图片格式、制作语义分割标签、划分 train/val/test。环境搭建创建独立环境安装 PyTorch、mmcv、mmsegmentation或直接用 Segformer 官方代码但更推荐 mmseg工程化做得更好。数据配置把自定义数据集注册进配置文件重点检查 image 和 label 的路径、类别数和 ignore_index。训练与验证设置合适的 batch size 和学习率开启 TensorBoard 或 wandb 看训练曲线及时止损。模型测试与导出用验证集或测试集图片进行推理输出彩色分割图必要时导出为 ONNX。这套流程不依赖特定业务场景遥感、工业、自动驾驶数据都能复用差别只在数据准备环节的标注工具和类别定义。2. 核心细节解析与实操要点2.1 训练环境与依赖版本必须按死一组能用的大版本环境问题往往是新手第一道坎。我踩过最难受的一次是 mmcv 和 mmsegmentation 版本不匹配报错全是 “cannot import name ‘BaseModule’ from ‘mmcv.runner’”查了很久才发现是 mmcv 装了 2.x而 mmseg 还是 0.x 的接口。后来我直接固定用下面这组版本组合基本没再翻车Python 3.8/3.9PyTorch 1.12.1 CUDA 11.6或者 PyTorch 1.13 CUDA 11.7mmcv-full 1.7.1mmsegmentation 0.30.0mmengine 0.7.4mmseg 0.30 会依赖它安装顺序很关键先建环境再装 PyTorch然后装 mmcv-full最后装 mmseg。很多同学先装 mmseg 再装 mmcv导致依赖解析把 mmcv 升到了 2.x报错就来了。conda create -n segformer python3.8 -y conda activate segformer pip install torch1.12.1cu116 torchvision0.13.1cu116 --extra-index-url https://download.pytorch.org/whl/cu116 pip install mmcv-full1.7.1 -f https://download.openmmlab.com/mmcv/dist/cu116/torch1.12/index.html pip install mmsegmentation0.30.0如果你显卡不支持 CUDA 11.6换成 cu113 的索引和 torch 1.10.2 也能跑但注意 torch 和 mmcv 编译版本的 cu 版本要一致否则会报 torch 的 CUDA 版本不匹配。2.2 数据标注与语义分割数据集的制作逻辑语义分割数据集最重要的就是 mask 图和 image 图要一一对应且 mask 必须是单通道或者三通道的“最像标注”格式。这里先说清楚很多人把 mask 做成彩色图每个类别涂一种颜色最后训练时发现 loss 不降原因就是类别编码没对齐。我用的规范是图片统一为 3 通道 RGB格式 JPG 或 PNG 都可以但最好统一。mask 必须是 PNG单通道灰度图像素值直接从 0 开始编号0 表示背景1 表示第一个类别2 表示第二个类别以此类推。如果标注工具导出的是彩色 mask例如 labelme 导出的 8 位彩色 PNG需要先做一个彩色到灰度索引的映射把 RGB 值映射为类别 id再另存为单通道 PNG。举个例子我的遥感数据有两个类别背景 0建筑物 1。那么 mask 中每个像素的像素值要么是 0 要么是 1。打开图片看应该是纯黑的背景上建筑物区域是淡白色像素值 1 显示为接近黑色的深灰而不是红色、绿色那种彩色块。这个转换可以用几行 Python 完成import numpy as np import cv2 color_map { # RGB - class (0, 0, 0): 0, # 背景 (255, 0, 0): 1, # 建筑物 } mask_color cv2.imread(label.png, cv2.IMREAD_COLOR) mask_color cv2.cvtColor(mask_color, cv2.COLOR_BGR2RGB) mask_gray np.zeros(mask_color.shape[:2], dtypenp.uint8) for rgb, class_id in color_map.items(): match np.all(mask_color np.array(rgb).reshape(1,1,3), axis-1) mask_gray[match] class_id cv2.imwrite(label_gray.png, mask_gray)转换完最好用 PIL 打开看一眼或者用下面代码统计像素值分布确认只有 0 和 1from PIL import Image import numpy as np img np.array(Image.open(label_gray.png)) print(np.unique(img))2.3 MMSeg 数据集格式必须遵守的目录结构MMSeg 支持自定义数据集但目录结构有约定。最省事的方式是模仿官方 Cityscapes 格式但更通用的是做一个标准的 “ImageMask” 结构data/ myseg/ images/ train/ img_001.png img_002.png ... val/ img_011.png masks/ train/ mask_001.png mask_002.png ... val/ mask_011.png注意 mask 的文件名和 image 的文件名不需要严格一致MMSeg 的默认数据集类会按顺序对应文件但为了排查方便我建议让 mask 文件名和 image 文件名保持一致比如img_001.png对应mask_001.png。这样你在做数据检查时一眼就能看出有没有错位。在配置文件里则需要显式指定数据集的类和路径。MMSeg 0.30 的配置是典型的 Python 文件我通常会复制segformer_mit-b0_512x512_160k_ade20k.py这类官方配置然后改成自己的数据集。2.4 训练配置文件的关键参数这是整个流程里最重要的一段。以我的遥感建筑物提取为例下面是精简后的 MMSeg 配置我逐行解释关键参数_base_ [ ../_base_/models/segformer_mit-b0.py, ../_base_/datasets/custom_seg.py, ../_base_/default_runtime.py, ../_base_/schedules/schedule_160k.py ] crop_size (512, 512) model dict( typeEncoderDecoder, backbonedict( typeMixVisionTransformer, in_channels3, embed_dims32, num_layers[2, 2, 2, 2], num_heads[1, 2, 5, 8], patch_sizes[7, 3, 3, 3], sr_ratios[8, 4, 2, 1], out_indices(0, 1, 2, 3), pretrainedpretrain/mit_b0.pth, init_cfgNone ), decode_headdict( typeSegformerHead, in_channels[32, 64, 160, 256], in_index[0, 1, 2, 3], channels256, dropout_ratio0.1, num_classes2, loss_decodedict( typeCrossEntropyLoss, use_sigmoidFalse, loss_weight1.0 ) ), train_cfgdict(), test_cfgdict(modewhole) )这里的num_classes2对应背景和建筑物。crop_size决定了训练时随机裁剪尺寸遥感图往往非常大比如 10000x10000直接整图训练显存不够所以训练时随机裁剪 512x512 的块。pretrainedpretrain/mit_b0.pth是指你下载好的预训练权重路径如果没有提前下载训练时会卡在下载环节建议离线提前下载到本地我后面会给出地址。test_cfgdict(modewhole)表示整图推理但实际如果图片太大也会爆显存我通常改成modeslide并设置crop_size和stride让模型滑动窗口推理。2.5 数据增强怎么配语义分割的数据增强比分类要谨慎因为几何变换会导致标签扭曲。MMSeg 在我的配置里常用这几个train_pipeline [ dict(typeLoadImageFromFile), dict(typeLoadAnnotations), dict(typeRandomResize, scale(512, 512), ratio_range(0.5, 2.0)), dict(typeRandomCrop, crop_sizecrop_size, cat_max_ratio0.75), dict(typeRandomFlip, prob0.5), dict(typePhotoMetricDistortion), dict(typePackSegInputs) ]cat_max_ratio0.75是防切出来的块全是某一个类比如全是背景限制单个类别占比不超过 75%避免模型被大块背景带偏。这个参数在目标物比较稀疏时特别有用。对于遥感影像我还会加RandomRotate90 度的倍数因为遥感图朝向可能是任意的但要注意rotate时标签也要同步旋转MMSeg 里已经有RandomRotate实现设置prob0.5即可。3. 实操过程与核心环节实现3.1 准备预训练权重和数据集跑通第一个 epoch先解决预训练权重。Segformer 官方在 GitHub 提供了 MIT-B0 到 B5 的 ImageNet-1K 预训练权重其中 MIT-B0 的下载地址一般在官方仓库的pretrained目录下如果访问不到也可以用 MMSeg 自带的下载脚本例如python tools/model_converters/convert_mit_backbone.py \ segformer_mit_b0.pth \ mit_b0.pth下载权重后放到pretrain/目录下在配置里写相对路径即可。数据集方面假设你已经有下面结构data/myseg/images/train/xxx.png data/myseg/masks/train/xxx.png ...然后在 MMSeg 的custom_seg.py数据集配置里dataset_type CustomDataset data_root data/myseg/ train_dataloader dict( batch_size8, num_workers4, persistent_workersTrue, samplerdict(typeDefaultSampler, shuffleTrue), datasetdict( typedataset_type, data_rootdata_root, data_prefixdict( img_pathimages/train, seg_map_pathmasks/train), pipelinetrain_pipeline) ) val_dataloader dict( batch_size1, num_workers4, datasetdict( typedataset_type, data_rootdata_root, data_prefixdict( img_pathimages/val, seg_map_pathmasks/val), pipelinetest_pipeline) )这里有个容易出错的点data_prefix的键名在 mmseg 0.30 和 mmseg1.x 不一样。0.30 用的是img_path和seg_map_path而 mmseg 1.x 改成了img和seg_map。我用的 0.30 是前者如果你装的是 1.x需要对应改。跑第一个 epoch 前先跑一下数据加载检查python tools/train.py configs/myseg/segformer_mit-b0_512x512_160k_myseg.py如果数据路径或标签有误通常会在data.loader阶段就报错比如 “FileNotFoundError”。这时候别急着改代码先检查路径前缀和文件名大小写Linux 上路径区分大小写Windows 不区分但容易埋坑。3.2 训练时长、学习率、batch size 的调参逻辑Segformer 原版是 160k iteration但那是大规模数据集的设置。自己数据集往往就几万张甚至几千张图不需要训练那么久。我的经验是数据量小于 2000 张训练 40k~80k iterations 足够。batch size 尽量在 8~16 之间显存不够就减小 crop_size 到 480 或 512。初始学习率建议1e-5到6e-5之间用poly衰减。预训练权重加载后学习率太大会破坏特征太小收敛太慢。我的调度器配置如下optim_wrapper dict( typeOptimWrapper, optimizerdict(typeAdamW, lr6e-5, betas(0.9, 0.999), weight_decay0.01), paramwise_cfgdict( custom_keys{ pos_block: dict(decay_mult0.), norm: dict(decay_mult0.), head: dict(lr_mult10.) } ) ) train_cfg dict(typeIterBasedTrainLoop, max_iters40000, val_interval2000)paramwise_cfg里我把 head 的学习率乘了 10因为解码器是随机初始化的需要更大步长去学backbone 预训练过用相对小的学习率去微调就好。val_interval2000表示每 2000 次迭代验证一次这样能及时看到 mIoU 是不是在涨。如果训练时 loss 在前期震荡很厉害可以把 lr 从 6e-5 降到 3e-5 甚至 1e-5。有些遥感影像颜色差异很小头几个 epoch 的 loss 掉得慢是正常的别急着停先跑 5k iteration 看趋势。3.3 训练过程中的监控与可视化训练太枯燥务必看曲线。MMSeg 默认支持 TensorBoard在启动命令里加--work-dir work_dirs/segformer_mit-b0_myseg后日志会写在 work dir 下然后tensorboard --logdir work_dirs/segformer_mit-b0_myseg浏览器打开http://localhost:6006主要看两个指标loss和mIoU。正常情况下 loss 是阶梯式下降mIoU 缓慢上升。如果 loss 一直不降优先怀疑数据标签问题而不是模型问题。我也习惯每 1000 次迭代保存一次 checkpoint这样一旦训练中断可以从最近的点恢复而不是从头再来。配置里这样加default_hooks dict( checkpointdict(typeCheckpointHook, interval2000, save_bestmIoU, rulegreater) )save_bestmIoU会自动保存验证集上 mIoU 最高的权重文件名带best标记这个权重用于后续测试。3.4 测试与推理加载最佳权重输出分割图训练完以后用验证集评估python tools/test.py configs/myseg/segformer_mit-b0_512x512_160k_myseg.py \ work_dirs/segformer_mit-b0_myseg/best_mIoU_iter_40000.pth \ --show-dir results--show-dir会保存可视化的分割图对比原图和预测效果。单独推理单张图可以写个小脚本from mmseg.apis import init_model, inference_model import cv2 model init_model(configs/myseg/segformer_mit-b0_512x512_160k_myseg.py, work_dirs/segformer_mit-b0_myseg/best_mIoU_iter_40000.pth, devicecuda:0) img cv2.imread(test.png)[..., ::-1] # BGR - RGB result inference_model(model, img) pred result.pred_sem_seg.data[0].cpu().numpy() # pred 是 HxW 的整数数组对语义分割网络而言pred的每个像素值是类别 id。想保存成彩色图就定义一个调色板把类别 id 映射成 BGR 颜色再cv2.imwrite。比如建筑物类别是 1我让它在输出图上显示为红色palette {0: (0, 0, 0), 1: (0, 0, 255)} # BGR vis np.zeros((*pred.shape, 3), dtypenp.uint8) for cls, color in palette.items(): vis[pred cls] color cv2.imwrite(result.jpg, vis)4. 常见问题与排查技巧实录4.1 类别数量不一致和数据不一致的原因这个是最常见的问题。报错信息往往是AssertionError: The num_classes (2) in head does not match the classes (3) in dataset出现这个的原因很简单——你标签里实际出现的类别数和配置文件里写的num_classes不一致。比如你的 mask 里既有 0、1还有不小心填出来的 2那即使你的本意是 2 类模型读到 3 个类别也照样报错。解决思路是先统计标签类别import numpy as np from PIL import Image import glob for f in glob.glob(data/myseg/masks/train/*.png): arr np.array(Image.open(f)) print(f, np.unique(arr))如果发现意外类别就需要重新处理标签。这类“数据不一致”问题排查起来极其耗时所以我在制作标签时就用脚本强制校验目标类别集合必须等于{0,1}一旦出现其他值直接打印文件路径并中断。另外还有一种数据不一致图片尺寸不一致。MMSeg 的RandomCrop能处理尺寸不同的图但如果你把尺寸差异特别大的硬塞到一个 batch 里且没有做RandomResize会报 stack 维度错误。我一般统一把所有训练图先 resize 到 1024x1024 或者 2048x2048再做训练这样省心。4.2 显存不足OOM的解决办法我自己第一次跑遥感大图时整图滑动窗口推理直接爆了 11G 显存。解决方案优先级如下减小crop_size从 512 改到 480 或 448显存占用同比下降约 12%~20%。减小 batch size从 8 改成 4必要时改成 2。开启梯度累积optim_wrapper dict( optimizerdict(typeAdamW, lr6e-5), accumulative_counts4 )这个设置相当于梯度积累 4 次再更新等效 batch size 为 8但显存占用和 batch size2 一样。使用混合精度训练AMPtrain_cfg dict( typeIterBasedTrainLoop, max_iters40000, val_interval2000, fp16dict(loss_scaledynamic) )注意 0.30 版本的 AMP 配置可能稍有出入但总思路是一致的。混合精度能省不少显存速度也有提升不过某些老旧显卡可能不太稳定如果出现 loss 突然 NaN就关掉 AMP 试试。4.3 精度不够、边缘粗糙、小目标丢失这里说点实在的。如果你的 mIoU 卡在某个值不动不要只改模型结构先从这三个方向排查标签边缘是否准确很多标注工具自动生成的边界非常粗如果标签本身边缘是“锯齿”的Segformer 再强也学不出平滑边界。建议把标签做一次 2~3 像素的腐蚀或者用标注工具的精修打磨。是否缺少多尺度推理测试时用modeslide配合多种 scale 做推理融合比如 [0.5, 0.75, 1.0, 1.25, 1.5]每个尺度输出概率图取平均后 argmax。这种方法能明显提高小目标召回率。类别不平衡是否严重如果只有 5% 像素是建筑物那么即使全预测为背景也能有很高准确率但 mIoU 就会很差。解决方法是给少样本类别更高的 loss 权重loss_decodedict( typeCrossEntropyLoss, use_sigmoidFalse, class_weight[0.3, 3.0], loss_weight1.0 )class_weight的顺序和类别 id 对应背景 0 权重 0.3建筑物 1 权重 3.0强制模型更关注少的那一类。4.4 从训练到部署的导出小坑训练后的模型要部署时很多人会直接转 ONNX。Segformer 转 ONNX 有两个常见问题一是它将GridSample算子用于特征重采样映射不标准二是动态输入尺寸不稳定。我在实际部署时是固定输入尺寸如 512x512并打开 opset 13用下面命令导出python tools/deployment/pytorch2onnx.py \ configs/myseg/segformer_mit-b0_512x512_160k_myseg.py \ work_dirs/best_mIoU_iter_40000.pth \ --output-file segformer.onnx \ --input-img test.png \ --shape 512 512导出后建议用 ONNX Runtime 跑一遍同一张图对比 PyTorch 输出差异如果像素级 argmax 结果基本一致就说明算子转换没问题。不同版本的 mmdeploy 坑还不一样但 0.30 配合这个固定 shape 的导出 route 是能用的。若你的生产环境是 GPU直接用 TensorRT 精度更好不过配置复杂度会更高新手可以先从 ONNX Runtime CPU 跑通全流程。4.5 小技巧速查表我把平时最常用到的避坑技巧整理成了一张表方便你直接对照现象原因快速解法训练 loss 不降标签类别编码错误或类别数配置错误检查 mask 的np.unique确认类别和配置一致mIoU 低但 loss 不高类别严重不平衡调整class_weight或对少数类做过采样边缘锯齿明显标签边缘粗糙精修标签或后处理时对输出概率图做 CRF验证内存爆掉整图推理显存不够test_cfg设置modeslide裁剪和步长都设 512恢复训练时 loss 变大学习率重新按全局 step 计算用load_from只加载权重不加载优化器状态或保留完整 checkpoint 恢复预训练权重下载失败网络访问问题手动下载到pretrain/目录修改配置pretrained为本地路径转换 ONNX 后结果全黑输入图片 BGR 与 RGB 顺序不一致确保推理时通道顺序和训练时一致通常为 RGB5. 一些经验之谈我是从遥感影像切图任务开始用 Segformer 的中间也折腾过 Deeplabv3、U-Net 和点云语义分割系列。整体感受是Segformer 在中小数据集上更容易得到一个“还不错”的结果但前提是数据和配置不出低级错误。有几个我后来反复强调的习惯这里再啰嗦一次第一所有数据必须跑一遍像素统计别相信眼睛看到的彩色图第二配置文件的num_classes一定要从标签里统计出来而不是拍脑袋第三每 2000 iteration 保存一次 checkpoint并且一定开 TensorBoard不能在黑盒状态下训练第四遇到不收敛先降学习率别一上来就换模型。另外Segformer 也不是万能的。如果你的业务对视实时性要求极高比如嵌入式端 30FPSB0 也未必跑得动可能要换轻量化的 MobileNet 系列二值化方案如果你的目标是超高精度遥感变化检测可能还要配合基于区域的方法。但作为通用语义分割 baselineSegformer 系列是我目前最推荐起步的模型之一。这次先写到这里关于数据制作后面我可以单独再聊一聊怎么用标注工具半自动提升标注效率以及怎样做多尺度后处理让 mIoU 再上一个台阶。希望这篇文章能帮你少走点弯路。