前段时间终于把SegFormer在4090上完整跑通了整个过程从驱动、CUDA、conda环境到预训练权重下载、官方仓库代码修改最后在一张遥感图上看到了干净的分割结果。SegFormer是NVIDIA开源的一套基于Transformer的语义分割方案核心是用分层Transformer编码器MiTMix Transformer加一个极轻量的全MLP解码器在性能和开销之间平衡得非常好也是很多做遥感解译、工业质检、街景语义理解的人首选的基础模型之一。这篇文章把我这次从零到一的完整过程整理了出来包含版本选型、环境配置、权重下载、代码改动、训练验证和测试评估。适合两类人一是想在4090这类消费级显卡上跑通SegFormer做实验的同学二是想用SegFormer做遥感或其他实际分割任务、需要快速落地训练和推理的工程师。照着走基本能让你少踩一半的坑。1. 项目整体设计与思路拆解1.1 为什么选4090 SegFormer这个组合4090看起来是个“消费级”显卡24GB显存却已经超过了上一代专业卡TITAN RTX而且Ada Lovelace架构对Transformer类模型的支持明显比30系更好TF32算力相当可观。SegFormer虽然带Transformer的字眼但它和ViT、Swin这类模型有个关键区别编码器使用分层金字塔结构解码器完全没有注意力模块所以整体参数量和显存占用比同等精度的Swin系列低不少。用4090跑SegFormer属于典型的“大炮打麻雀但打得非常舒服”的组合——你甚至可以在显存里塞下大batch或者更高分辨率输入。另外一个现实原因是SegFormer官方开放的代码仓库是基于MMSegmentation 0.x版本的这个版本生态在新卡上配置多多少少会碰到兼容性问题。4090太新而老代码、老工具链往往还是为30系甚至20系准备的。所以这个项目最核心的工作不只是“跑通模型”而是“让老代码在新硬件上稳定跑起来”。1.2 从零到一的完整技术链路这次项目整体上可以拆成一条清晰的链路我建议你也按这个顺序推进不要跳步安装NVIDIA驱动确保nvidia-smi能正常识别4090。准备CUDA Toolkit确保编译工具链可用。创建conda独立环境安装Python 3.9和匹配的PyTorch。安装mmcv-full和mmsegmentation版本必须严格匹配PyTorch和CUDA。克隆SegFormer官方仓库和对应版本的MMSegmentation仓库理清两者关系。下载MiT系列预训练权重配置到模型初始化路径。修改数据集路径、类别数、批量大小等配置。依次执行单图推理、训练、测试评估。这条链路里最容易出问题的集中在第4步和第7步。版本匹配直接决定了能不能import成功配置修改决定了训练能不能收敛、测试指标是不是合理。1.3 版本选型的核心决策提到版本我先把我这次最终使用的组合放在这里避免你被各种散落文档绕晕系统Ubuntu 24.04 LTS显卡驱动NVIDIA 550系列CUDA Toolkit11.8实际用PyTorch自带的Runtime也能跑但为了方便编译工具链建议装Python3.9PyTorch2.0.1 cu118mmcv-full1.7.2预编译包对应torch 2.0.1 / cu118mmsegmentation0.30.0SegFormer官方仓库master分支为什么不用最新的mmsegmentation 1.x主要原因在于SegFormer官方仓库的config格式和demo脚本都基于0.x新版mmseg的数据结构、config字段和工具脚本变化很大直接用新版反而要改更多代码。先跑通官方链路拿到可信的baseline再考虑迁移到新版本这才是稳妥路线。2. 环境配置驱动、CUDA与PyTorch工具链2.1 先装驱动nvidia-smi要能看到卡如果你用的是Ubuntu 24.04安装驱动其实比老版本简单很多。系统装好后只有一个基础的nouveau驱动跑深度学习是远远不够的。先查一下推荐驱动版本ubuntu-drivers devices输出里会列出当前机器上可用的NVIDIA驱动通常带有recommended标注的那个就是系统推荐版本。对于4090推荐版本大概率是550系列或者更高直接安装sudo apt update sudo apt install nvidia-driver-550装完后重启执行nvidia-smi。正常情况下你会看到类似这样的输出----------------------------------------------------------------------------- | NVIDIA-SMI 550.xx Driver Version: 550.xx CUDA Version: 12.4 | |--------------------------------------------------------------------------- | GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. | || | 0 NVIDIA GeForce RTX 4090 On | 00000000:0X:00.0 Off | 0% |这里有一个非常普遍的误解需要提前纠正nvidia-smi右上角显示的CUDA Version并不是你系统里安装的CUDA版本而是当前驱动“最多能支持到哪个CUDA版本”。也就是说驱动只要显示CUDA 12.4那么你后面装CUDA 11.8或12.1的Toolkit都是没有问题的驱动会向下兼容。2.2 CUDA Toolkit与驱动的真实关系为什么明明PyTorch自带了CUDA Runtime还是建议装一个系统级CUDA Toolkit因为mmcv-full这类库在安装或编译时需要通过nvcc获取CUDA编译器信息如果你系统里没有nvcc就只能依赖别人预编译的wheel包。预编译wheel一但找不到对应torch版本就必须自己编译这时候没有nvcc就直接卡死了。CUDA Toolkit 11.8的安装我建议用runfile方式不要用apt避免它顺手把驱动也替换掉wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run sudo sh cuda_11.8.0_520.61.05_linux.run安装界面里记得取消勾选Driver只装CUDA Toolkit那一部分。装完后配置环境变量export PATH/usr/local/cuda-11.8/bin:$PATH export LD_LIBRARY_PATH/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH把这两行追加到~/.bashrc然后source ~/.bashrc。验证是否生效nvcc -V能看到release 11.8, V11.8.89之类的输出就说明工具链就绪。2.3 conda环境里装PyTorch并验证GPUPyTorch的安装特别强调版本匹配。我这次用的是torch 2.0.1对应CUDA 11.8的预编译包conda create -n segformer python3.9 -y conda activate segformer pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118装完立刻验证能不能看到4090python -c import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))正常输出2.0.1cu118 True NVIDIA GeForce RTX 4090如果torch.cuda.is_available()返回False多半是驱动有问题直接回到2.1节检查nvidia-smi。这一步通过之后后面所有环境问题都不需要再怀疑GPU了。2.4 mmcv-full、mmsegmentation与官方仓库的匹配这是整个环境配置里最容易翻车的一步。mmcv的编译强依赖PyTorch和CUDA的精确版本所以OpenMMLab官方提供了一组预编译索引我们直接按版本拉取pip install mmcv-full1.7.2 -f https://download.openmmlab.com/mmcv/dist/cu118/torch200/index.html命令里的cu118对应CUDA 11.8torch200对应PyTorch 2.0.x。你只要确认这两个都对得上就能安装成功。安装完检查python -c import mmcv; print(mmcv.__version__, mmcv.is_compiled_with_cuda())输出1.7.2 True就说明mmcv的CUDA算子已经编译并加载成功。这个检查很重要因为mmcv这类库如果只是纯Python模块而没有CUDA扩展很多分割功能会悄悄退回CPU或者直接报错。之后再装mmsegmentationpip install mmsegmentation0.30.0很多教程喜欢让你git cloneMMSegmentation以后用pip install -e .安装但我实测在0.30.0这个版本上直接pip安装就已经够用因为后面我们还需要clone一份mmsegmentation仓库来使用它的tools脚本两者互不冲突。3. 预训练权重下载与安置3.1 需要哪些权重文件SegFormer的模型族谱可以分成两部分。第一部分是骨干网络MiT的预训练权重有B0到B5六个规格这些是在ImageNet-1k上预训练出来的文件名大概是mit_b0.pth、mit_b1.pth这种第二部分是针对具体分割数据集的完整模型权重比如segformer.b0.512x512.ade.160k.pth是在ADE20K上训练好的完整权重。这里要提醒你一个新手最容易踩的坑如果你想做的是自定义数据集训练你需要的是前一类也就是MiT的骨干预训练权重它只负责初始化backbone部分。如果你只是想跑官方demo看效果直接下载第二类对应数据集的完整权重即可。3.2 下载方法与镜像加速官方权重主要托管在Google Drive国内网络环境下载经常掉线重试非常痛苦。比较推荐的方式是从HuggingFace拉取社区用户已经把完整权重同步上去了。如果你访问HuggingFace本身也比较费劲可以先设置环境变量切换到国内可用镜像export HF_ENDPOINThttps://hf-mirror.com然后直接用huggingface-cli或wget下载。例如拉取MiT-B0的backbone权重wget https://hf-mirror.com/zhiqiangc/SegFormer/resolve/main/mit_b0.pth下载完建议用du -h mit_b0.pth看一眼大小B0大概在30MB左右B5能到500MB以上。如果文件明显偏小说明下载中途断了赶紧重新下载别等到训练时报EOFError再回头查。3.3 权重验证与路径配置拿到.pth文件后可以顺手用torch打开看一眼确认里面没问题python -c ckpt__import__(torch).load(mit_b0.pth); print(type(ckpt)); print(list(ckpt.keys())[:5])SegFormer的backbone权重通常是一个纯state_dict或者包含state_dict键的字典keys里面会出现backbone.patch_embed1.proj.weight这类字段。看完确认结构正常再把它放到固定目录。我习惯在项目根目录建一个pretrained/文件夹统一存放SegFormer/ ├── local_configs/ ├── pretrained/ │ └── mit_b0.pth ├── data/ └── work_dirs/后面改config时pretrained字段直接写相对路径或绝对路径都可以但绝对路径最保险尤其是当你同时操作多个conda环境或从不同工作目录启动脚本时相对路径很容易找错。4. 代码修改与数据集适配4.1 把官方仓库和MMSegmentation的关系理顺SegFormer官方仓库本身并不自带训练和测试脚本它的核心是local_configs里的模型配置文件。真正执行训练时还是要依赖MMSegmentation仓库里的tools/train.py、tools/test.py。所以推荐的做法是同时clone两个仓库然后把SegFormer的config复制到MMSegmentation的configs目录下或者直接用绝对路径引用比较灵活。git clone https://github.com/NVlabs/SegFormer.git git clone -b v0.30.0 https://github.com/open-mmlab/mmsegmentation.git如果你用的是SSH把https://换成gitgithub.com:也是一样的。完成之后在MMSegmentation目录下创建一个configs/segformer目录把SegFormer/local_configs/segformer/下的内容链接过去ln -s /path/to/SegFormer/local_configs/segformer /path/to/mmsegmentation/configs/segformer这样configs/segformer/下就能看到B0、B1等一系列子目录训练脚本也能通过标准路径加载配置。4.2 核心配置文件逐段拆解以segformer.b0.512x512.ade.160k.py为例这个文件本身其实是在继承基础配置的基础上做微调。真正核心的部分在它的基类里比如base/segformer_b0-512x512-160k_ade20k.py。打开后你主要看这几个字段model dict( typeEncoderDecoder, pretrainedpretrained/mit_b0.pth, # 换成你的backbone权重路径 backbonedict(typeMixVisionTransformer, ...), decode_headdict( typeSegformerHead, in_channels[32, 64, 160, 256], in_index[0, 1, 2, 3], channels256, dropout_ratio0.1, num_classes150, # 改成你自己的类别数 ... ) )这里有两个核心修改点。第一pretrained字段如果是自定义数据集请务必将它指向你下载好的MiT权重如果你想从头训练但没有任何预训练也可以直接删除这个字段但收敛速度会明显变慢最终精度也大概率不理想。第二num_classes官方默认150是ADE20K的类数Cityscapes是19类你的遥感数据如果只有5类就改成5。4.3 组织数据集以遥感图为例SegFormer的数据集格式遵循MMSegmentation的通用约定图片目录和标注目录分离。以遥感单类别分割为例目录结构这样组织data/ └── remote/ ├── img_dir/ │ ├── train/ │ │ ├── 001.png │ │ └── 002.png │ └── val/ │ ├── 003.png │ └── 004.png └── ann_dir/ ├── train/ │ ├── 001_mask.png │ └── 002_mask.png └── val/ ├── 003_mask.png └── 004_mask.png图片和标注的文件名必须一一对应标注是单通道PNG背景像素值为0目标类别像素值为1、2、3这样从1开始递增255作为ignore_index。然后回到config里修改数据配置data dict( samples_per_gpu2, workers_per_gpu4, traindict( typeCustomDataset, data_rootdata/remote/, img_dirimg_dir/train, ann_dirann_dir/train, num_classes1, pipelinetrain_pipeline), valdict( typeCustomDataset, data_rootdata/remote/, img_dirimg_dir/val, ann_dirann_dir/val, num_classes1, pipelinetest_pipeline) )这里要注意CustomDataset默认会从标注文件里读取label来推断类别数但我还是建议显式写上num_classes和classes字段避免后续评估时类别数对不上。如果标注像素值不是从0开始而是在1、2、3之间跳变训练时很容易出现loss不下降或者mIoU恒为0的情况遇到这种情况优先检查标注文件里的像素值到底有哪些。4.4 单卡训练和demo脚本的几个必改点官方config是按多卡训练设计的samples_per_gpu虽然是2但通常配合nnodes1、nproc_per_node8来跑总batch size是16。在4090上跑单卡你必须手动控制samples_per_gpu否则8卡调度代码会直接卡住或报错。启动训练时用--cfg-options覆盖即可不需要改config文件python tools/train.py configs/segformer/B0/segformer.b0.512x512.ade.160k.py \ --work-dir work_dirs/b0_remote \ --cfg-options data.samples_per_gpu2 data.workers_per_gpu4如果需要顺便修改训练分辨率也可以在--cfg-options里加上data.train.pipeline.2.img_scale(512,512)这类参数但更稳妥的做法是直接改config文件里的train_pipeline。demo推理脚本相对简单mmsegmentation/demo/image_demo.py是现成的直接指定一张图片、一个config、一个checkpoint就能跑。但官方脚本默认从mmseg仓库的configs目录找目标文件所以你引用外部config时最好用绝对路径。5. 运行、验证与测试全流程5.1 用预训练权重快速跑出一张分割图先跑通推理这是验证整个环境最直接的方式。我建议下载一个官方在ADE20K上训练好的完整权重比如segformer.b0.512x512.ade.160k.pth然后执行python demo/image_demo.py /path/to/test.jpg \ /path/to/SegFormer/local_configs/segformer/B0/segformer.b0.512x512.ade.160k.py \ /path/to/checkpoint/segformer.b0.512x512.ade.160k.pth \ --device cuda --palette ade20k --out-file result_ade.png执行后你会看到脚本输出一个result_ade.png不同类别被渲染成不同颜色。这一步通过说明驱动、CUDA、PyTorch、mmcv、mmseg整个环境已经全部打通。如果你需要用自己训练的遥感权重跑同样的推理把config换成你的训练配置checkpoint换成你训练产出的latest.pth或iter_xxx.pth模型定义的类别数、解码器结构会和训练时保持一致不需要担心推理代码层面的适配。5.2 小数据集训练日志怎么看、显存怎么控跑通demo后建议先用一个小数据集训练一个很短周期验证训练管线没有问题。启动训练后日志大约是这个样子2025-04-01 12:00:01,234 - INFO - Exp name: segformer.b0.512x512.ade.160k 2025-04-01 12:00:02,101 - INFO - Iter [100/160000] lr: 1.0e-4, eta: 4:28:20, time: 0.345, data_time: 0.012, memory: 3218, loss: 2.5831 2025-04-01 12:00:07,520 - INFO - Iter [200/160000] lr: 1.0e-4, eta: 4:20:15, time: 0.323, data_time: 0.010, memory: 3220, loss: 2.1024这里值得关注的三个信息time是单次迭代耗时4090上跑512x512的B0模型单卡迭代时间在0.3s左右算是正常memory是当前batch的显存占用单位是MB如果接近24000就要小心OOMloss应该随迭代逐步下降如果连续几百步都不动马上停下来查数据集标注和类别配置。显存控制有几个常用手段按优先级排序降低samples_per_gpu从2降到1通常能省出接近一半的显存。减小输入分辨率把512x512改成384x384显存下降非常明显。打开混合精度训练在config里添加fp16dict(loss_scaledynamic)实测能省20%到30%显存。5.3 测试评估mIoU从哪里来训练完成后用官方tools/test.py进行评估。我这次用的命令是python tools/test.py configs/segformer/B0/segformer.b0.512x512.ade.160k.py \ work_dirs/b0_remote/latest.pth --eval mIoU如果验证集配置正确输出会是一张漂亮的表格-------------------------- | Scope | mIoU | mAcc | -------------------------- | global | 42.31 | 53.12 | --------------------------mIoU是所有类别IoU的平均值mAcc是平均类别准确率这是语义分割最常用的两个指标。如果你的验证集是多类别建议再加上--eval mIoU mAcc aAcc一次拿到更多指标。5.4 可视化与结果落盘除了指标你大概率还想看看模型到底预测成了什么样。tools/test.py支持--show-dir参数会在指定目录落盘可视化结果python tools/test.py config.py latest.pth --eval mIoU --show-dir vis_output它会在vis_output/下生成三列对比图原始图、真值图、预测图。这个可视化结果对找bad case非常有帮助比如某类目标频繁被漏检直接看预测图就能定位是训练样本不均衡还是感受野不够导致小目标丢失。6. 常见问题与排查技巧实录6.1 六个高频报错与解法我这次配环境过程中遇到的坑和对应解法整理成了下面这个表建议收藏备用现象原因解决办法torch.cuda.is_available()为 FalseNVIDIA驱动没装好或版本过低执行nvidia-smi确认驱动正常重新安装550及以上版本驱动No module named mmcv._extmmcv的CUDA扩展没编译成功用mim install mmcv-full1.7.2重装或从OpenMMLab官网预编译索引安装对应torch/cuda版本undefined symbol: _ZN2at6detail...mmcv与PyTorch版本不匹配严格按torch 2.0.1cu118对应mmcv-full 1.7.2的组合安装不要混装CUDA out of memorybatch size或分辨率过大降低samples_per_gpu、降低输入尺寸、开启fp16训练loss不降且mIoU恒为0数据集类别配置错误或mask像素值从1开始检查ann_dir中的mask取值应包含0且最大类数不超过num_classes让像素值连续从0开始NVIDIA-SMI has failed because it couldnt communicate with the NVIDIA driver驱动内核模块加载失败重启系统如果之前装过runfile驱动需要先清理旧驱动再装apt版6.2 版本兼容问题的排查思路如果你遇到的是上面表格之外的版本类报错我的排查思路一般是三步。第一步确认PyTorch的CUDA版本和编译用的CUDA Toolkit版本一致可以在终端分别执行python -c import torch; print(torch.version.cuda) nvcc --version这两者完全一致才能最大程度避免问题。第二步确认mmcv和mmseg版本配对mmseg 0.30.0对应的mmcv-full最低版本是1.5.3但1.7.2实测工作良好。第三步如果还是报错优先考虑隔离环境从零重装不要试图在现有环境里反复升级降级那样往往会把依赖树搞得一团糟。6.3 在4090上省显存和加速的几条经验最后分享几条在4090上实测有效的经验。第一4090的Tensor Core对半精度计算非常友好如果你的任务是比较重的数据集强烈建议全程开FP16训练速度能提升30%以上显存占用明显下降。第二单卡训练时不要直接拿官方多卡学习率用官方在batch size 16下设置的base_lr是0.00006你单卡batch size如果是2按线性缩放大概用0.00006 * 2 / 16 0.0000075显然太小实践中可以直接上调一个数量级到0.0001配合warmup基本稳定。第三workers_per_gpu不要盲调太高4个就够了太高反而会因CPU换页开销拖慢训练。我自己在跑SegFormer的过程中最大的体会是这类型老代码的配置难度其实不在于某个单一环节而在于版本链条上任何一环脱节都会让你在半夜对着一个莫名其妙的报错发呆。所以动手前把版本清单写清楚每个环节安装完立刻验证是省时间的关键。另外一个小技巧在你第一次启动训练前先随便找一张验证集小图跑一次demo推理确认模型能正常前向传播再启动长时间训练能帮你提前发现90%的配置问题。