FEATURED · 精选文章

RTX 4090跑通SegFormer完整指南:环境配置、踩坑实录与代码适配

发布时间 / 2026/9/20 5:13:08
来源 / 创域科博编辑部
栏目 / 资讯中心
RTX 4090跑通SegFormer完整指南:环境配置、踩坑实录与代码适配 机器是上周到的4090插上去第一件事就是想把官方SegFormer复现出来结果栽在环境上一个下午。你去看NVIDIA/SegFormer仓库里的environment.yaml里面还写着pytorch 1.8、cuda 10.2这在4090上根本起不来。如果你也是刚换了4090、又想在官方代码基础上跑分割实验这篇就是把你从“装完驱动不知道下一步”带到“单图验证通过、测试脚本能跑”的完整过程包含预训练权重下载、conda环境组合、代码改动点以及几个我实际踩到、搜索引擎也救不了的坑。1. 为什么在4090上跑官方SegFormer第一件事是忘掉官方environment.yaml1.1 4090的sm_89与老版本PyTorch/CUDA的错配RTX 4090用的是Ada Lovelace架构计算能力是sm_89。这里有个非常实际的问题CUDA扩展比如mmcv里的自定义算子在编译时会针对特定算力生成机器码。PyTorch 1.8时代默认只编译到sm_50~sm_80附近根本没有sm_89的SASS机器码。显卡加载老内核时只能尝试用PTX JIT即时编译一旦算子复杂一点就极其容易崩或者干脆直接报“no kernel image available”。这还不是全部。老版本CUDA Toolkit官方环境里写的10.2在Linux下搭配新版NVIDIA驱动经常出现驱动版本过高、与旧CUDA runtime不兼容的告警有些场景直接起不来。所以一句话在4090上凡是“pytorch 1.8 cuda 10.2 mmcv-full 1.3.x”这种组合理论上有机会实操就是血压升高。1.2 官方仓库的版本依赖到底卡在哪SegFormer官方仓库的核心依赖其实不多真正麻烦的是三个包的版本联动组件官方环境写的版本4090上建议版本原因Python3.73.83.7太老很多新版编译链不再提供wheelPyTorch1.8.01.13.1cu1171.13是mmcv-full 1.7.2官方预编译wheel覆盖到位的版本CUDA Toolkit10.211.7 或 11.8驱动层和编译层必须支持sm_89mmcv-full1.3.91.7.21.7.2对CUDA 11.x适配成熟且提供官方预编译包mmsegmentation0.13.00.20.2与mmcv-full 1.7.2搭配稳妥能直接读官方repo的config为什么不是PyTorch 2.0因为mmcv-full 1.7.2没有针对torch 2.0的官方预编译包需要自己从源码编而编译的变数远大于预编译。如果你硬要torch 2.0就得走mmcv 2.0.x mmsegmentation 1.0.x但SegFormer仓库里那一堆老式config格式在mmseg 1.0里要做迁移改动不划算。老老实实1.13.1所有精力留在模型上。2. 从NVIDIA驱动到conda环境一套在4090上实测通过的配置组合2.1 显卡驱动与CUDA Toolkit的版本对应关系先装驱动。4090刚上市那会儿装驱动是个坎现在驱动版本已经很成熟直接用535或545系列就行。确认驱动装好用nvidia-smi看输出注意CUDA Version那一栏显示的是“驱动支持的最高CUDA版本”它和你在conda里装的CUDA Toolkit是两码事。很多新手在这里被绕晕驱动显示的CUDA 12.2不代表你的PyTorch就能用CUDA 12.2PyTorch通过自己的CUDA runtime运行你只需要保证驱动版本足够新就行。我的建议是驱动装完直接用conda装PyTorch时自带的CUDA 11.7 runtime完全不用单独装CUDA Toolkit。cuDNN也一样PyTorch官方wheel里已经带了不再需要系统级安装。2.2 conda创建虚拟环境与PyTorch安装创建环境这一步命名随意但Python版本一定锁3.8conda create -n segformer python3.8 -y conda activate segformer接着装PyTorch。网上很多教程让你pip install torch torchvision这在4090上有风险因为默认源的包可能不带CUDA算子必须指定cu117的index-urlpip install torch1.13.1 torchvision0.14.1 --index-url https://download.pytorch.org/whl/cu117如果下载慢可以带上--proxy之类的手段或者挑网络空闲时段。装完务必验证一下CUDA是否可用python -c import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))看到1.13.1cu117 True NVIDIA GeForce RTX 4090这一步才算过。2.3 mmcv-full安装编译还是预编译包这里最容易翻车mmcv-full是整个环境里最容易翻车的组件因为它包含大量CUDA自定义算子比如deformable attention、carafe、sync_bn等。安装前建议先把ninja和gcc装好避免中途失败后无从下手pip install ninja sudo apt-get install gcc g # Ubuntu/Debian其他系统用对应包管理器有两条路可以走路径A用openmim装预编译包推荐优先试pip install openmim mim install mmcv-full1.7.2openmim会根据你当前环境里的PyTorch和CUDA版本自动选择合适的wheel。如果你的PyTorch是1.13.1cu117它大概率能拉到预编译版本直接几分钟装完。路径B源码编译预编译拉不到时走这条git clone https://github.com/open-mmlab/mmcv.git cd mmcv git checkout v1.7.2 export FORCE_CUDA1 export TORCH_CUDA_ARCH_LIST8.9 export MMCV_WITH_OPS1 pip install -e .这里的TORCH_CUDA_ARCH_LIST8.9是4090的算力编号。不设这个变量的话编译时会尝试把能想到的所有算力全部编一遍耗时翻好几倍而且容易在老架构编译环节报错。只锁定8.9干净利落。装完mmcv-full接着装mmsegmentationpip install mmsegmentation0.20.2最后回到SegFormer官方仓库把仓库clone下来git clone https://github.com/NVIDIA/SegFormer.git cd SegFormer到这里环境部分基本齐了。建议现在就把环境整体保存一下比如记录一份requirements.txt因为接下来装权重、跑模型的过程中你随时可能因为某个操作把环境搞乱有个记录能快速恢复。3. 预训练权重下载、命名规律与本地化路径修正3.1 官方权重放哪、文件名与config的对应关系SegFormer的权重分为两段一是mix_transformer的backbone预训练权重二是整个模型在ADE20K等数据集上finetune后的完整权重。你实际测试时直接用完整权重不用管backbone那一段。官方仓库的README.md里每个模型都给了Google Drive和Hugging Face的下载链接。文件名规律很直观segformer.b0.512x512.ade.160k.pthB0模型输入512x512ADE20K训练160k iterationsegformer.b1.512x512.ade.160k.pth一直到segformer.b5.640x640.ade.160k.pth下载后统一放到一个目录比如SegFormer/pretrained/后续测试脚本会用到。这里特别说一句如果你的网络环境对海外下载不友好Hugging Face镜像站和Google Drive的第三方下载器都是替补方案。重点是下载完确认文件大小。ADE20K完整权重大概是几百MB级别B5会更大一些如果一个文件只有几十MB那基本是下载中断或下载了个错误页面别往下跑直接重新下载。3.2 权重路径写死导致“找不到文件”的解决办法官方config文件里model.backbone.init_cfg.checkpoint默认指向一个http地址也就是ImageNet预训练权重的路径。你本地跑测试时这个地址能不能访问都是问题而且也没必要再下那一段权重。测试的时候完整模型权重是通过命令行参数--checkpoint传进去的用的是mmseg里load_checkpoint的逻辑它不会去看config里那个init_cfg.checkpoint。所以如果你只是跑官方测试命令下载完整权重就够了不需要改config里的预训练路径。但如果你后续要finetune自己的数据集那就要改两个地方一是把init_cfg.checkpoint改成你本地backbone权重路径二是把model.decode_head.num_classes改成自己的类别数。很多人在这一步把pretrained和checkpoint搞混花几个小时排查“为什么没加载权重”——其实一个是训练初始化用一个是你加载进来评估/推理用属于两套机制。3.3 权重文件损坏的快速判断下载下来的pth文件如果损坏了跑测试时不会马上报“文件损坏”而是会报类似“size mismatch”或者unexpected key因为有些层的形状对不上或者加载到一半就EOFError。碰到这种错误第一反应不是去改代码而是先比对文件大小、重新下载。我习惯下载完先做个加载自检python -c import torch; ckpttorch.load(pretrained/segformer.b0.512x512.ade.160k.pth, map_locationcpu); print(type(ckpt), len(ckpt) if isinstance(ckpt, dict) else )能正常打印出dict结构和长度说明文件基本完整可以进入下一步。4. 代码适配让官方Config和测试脚本在新显卡上跑起来4.1 config里必须调整的采样与同步批量配置SegFormer官方repo的config文件放在local_configs/segformer/下以B0为例核心配置在segformer.b0.512x512.ade.160k.py里。对单卡4090测试来说默认的samples_per_gpu是1用不着改但有两个点需要留意。第一data.test的img_scale默认是(2048, 512)之类的大尺寸。测试时mmseg会用这个尺度做resize对单图推理来说没毛病但如果你拿的是自己随手拍的遥感图或街景图分辨率比例跟默认不一致输出结果会跟你预期差很多。建议先保持默认跑通再根据自己的图片尺寸调整。第二model.decode_head.loss_decode里的loss_weight和num_classes在ADE20K权重下是150类不需要动。如果换到自己数据集num_classes必须改成自己的类别数否则size mismatch立刻找上门。4.2 单图验证脚本不依赖数据集一张图看出流程是否通了官方repo的tools/test.py主要面向数据集评估对“我就想拿一张图试试”的场景并不友好。所以我自己写了一个极简推理脚本逻辑清晰方便你逐步验证环境、权重、预处理流水线是否全通。import torch import mmcv import numpy as np from mmseg.models import build_segmentor from mmseg.apis import inference_segmentor, init_segmentor config_file local_configs/segformer/B0/segformer.b0.512x512.ade.160k.py checkpoint_file pretrained/segformer.b0.512x512.ade.160k.pth # 模型初始化显式指定GPU model init_segmentor(config_file, checkpoint_file, devicecuda:0) # 读取任意图片输出是list里面是HxW的预测标签图 img_path test.jpg result inference_segmentor(model, img_path) # 保存可视化结果附件会叠加在原图上 model.show_result(img_path, result, out_filetest_seg.jpg, opacity0.6) # 如果想看各类别像素分布可以用np.bincount统计 labels, counts np.unique(result[0], return_countsTrue) for lab, cnt in zip(labels, counts): print(fclass {lab}: {cnt} pixels)这个脚本里init_segmentor会自己读取config、加载checkpoint、构建模型并加载权重不用手动load_state_dict避免了很多人踩的“权重键名对不上”问题。第一次跑如果没报错并且能在test_seg.jpg里看到清晰的语义分割色块就说明环境、权重、推理链路全部打通了。4.3 官方test.py跑数据集评估时需要注意的路径与参数格式想复现论文里的mIoU就要用官方数据集评估。以ADE20K为例先把数据集下载好目录结构按mmseg要求组织data/ ade/ ADEChallengeData2016/ annotations/ training/ validation/ images/ training/ validation/然后修改config中的data_root。官方config里写的是data_root data/ade/ADEChallengeData2016如果你的数据集在别的路径改这个地方别去改一堆data.train、data.val里的具体路径因为它们都基于data_root拼接。跑评估的命令python tools/test.py local_configs/segformer/B0/segformer.b0.512x512.ade.160k.py \ pretrained/segformer.b0.512x512.ade.160k.pth \ --eval mIoU如果你的显存或显存带宽足够可以把--eval后面加上mDice、mFscore一起评估。4090的24G显存跑B0验证集完全没压力。评估结果会在终端里逐类打印IoU最后汇总mIoU。我这个环境里B0在ADE20K val上复现到37%左右和论文报告的37.32%对得上说明环境和权重属实没问题。5. 运行与测试用一张图和验证集分别验证流程是否真的通了5.1 单图推理的完整输出与预期结果用第4节的脚本跑一张ADE20K里的室內图你会看到这样的输出class 0: 12345 pixels class 10: 4567 pixels class 88: 1024 pixels ...预测标签的数值范围是0到149ADE20K共150类这些值对应CLASSES元组里的具体类别名model.CLASSES可以直接查看。可视化输出test_seg.jpg里同一类别的区域会统一成一个色块边界轮廓清晰说明decoder部分工作正常。如果输出图是一团乱麻、色彩完全不对常见原因有两个一是预处理时图片被错误resize导致模型输入比例失衡二是权重加载不完整模型实际是随机初始化状态。排查顺序是先看init_segmentor有没有输出“load checkpoint from ...”的日志再检查图片路径的读取是否正常。5.2 验证集评估的耗时与显存参考用官方tools/test.py跑ADE20K val共2000张图在4090上B0模型大概需要2到3分钟B5模型在640x640输入下大约10分钟左右具体取决于显卡功耗限制和驱动版本。显存方面B0在batch size 1下大概占2.5GBB5占5GB左右24G显存远远够用。模型输入尺寸推理显存占用batch1val 2000张耗时参考ADE20K val mIoUSegFormer-B0512x512约2.5GB约2-3分钟约37.3%SegFormer-B2512x512约3.8GB约4分钟约45.6%SegFormer-B5640x640约5GB约10分钟约49.5%注这些耗时是在我的机器上的量级参考不同驱动版本和散热策略会有波动但有一个趋势是确定的SegFormer的高效性在4090上体现得非常明显B5也就这个量级。5.3 4090上特别值得关注的GPU利用率问题跑测试时盯着nvidia-smi看你可能会发现GPU利用率只有30%到60%显存也没吃满。这不是代码有问题而是SegFormer的解码器本身轻量batch size又只有1主力瓶颈在数据加载和预处理上GPU根本没有被喂饱。想榨干性能可以在DataLoader参数里调高num_workers减少CPU预处理拖后腿把测试的samples_per_gpu提到2或4提升GPU利用率打开cudnn.benchmark True让cuDNN针对当前输入尺寸做算法选择不过这些都是后期优化第一轮验证老老实实用batch1跑通别一上来就开优化出了问题会更难排查。6. 踩坑实录这几天我遇到过的报错和排查链路6.1 “No kernel image available”的完整追溯过程这个报错是我接手4090后遇到的第一道坎信息量很少就一句话。当时我的排查链路是nvidia-smi看驱动显示正常CUDA版本12.x。用torch.cuda.get_device_name(0)验证PyTorch能识别显卡识别到了。随便跑个torch.randn(100, 100).cuda() torch.randn(100, 100).cuda()正常。一旦跑SegFormer的decode head立刻崩报No kernel image available。到这里才意识到PyTorch基础张量操作能跑不等于所有CUDA算子都能跑。mmcv的自定义算子里有些编译目标是旧架构在4090上缺少对应SASS。解决方式就是把mmcv版本升级到我在第2节写的那套组合让所有算子都带上sm_89的机器码。这个报错给的经验是遇到CUDA相关崩溃先分清楚是PyTorch基础库的问题还是第三方扩展库的问题排查范围能缩小一半。6.2 mmcv版本兼容性引发的SegFormer head报错另一个高频报错是AttributeError: ConfigDict object has no attribute xxx或者TypeError: format() argument indent must be an int, not None这类报错十有八九是mmcv版本和mmsegmentation版本不匹配。mmcv 1.7.x的ConfigDict行为和旧版默认值有差异而mmseg 0.13.0的开发时间线更早两者接口对接时就容易出现奇葩报错。我当时的解决路径是先用pip list | grep mm检查版本组合发现装的是mmcv-full 1.7.2 mmsegmentation 0.13.0。查阅open-mmlab的版本对应表发现1.7.2对应的mmsegmentation应该在0.20.0以上。升级到mmsegmentation0.20.2报错消失。这就是为什么我整个“环境搭建”章节里反复强调版本组合。你单独看每一个包都没问题但它们凑在一起可能就有问题最好直接按经过验证的组合来。6.3 印象最深的两个“小时级”坑apex与triton这两个坑都会在编译或启动时冒出来但本质完全不同。第一个是apex。很多老教程会在装SegFormer时让你顺手装一下NVIDIA的apex优化库说能加速。但这个库在4090上遇到新版PyTorch经常在编译阶段卡死而且SegFormer官方代码根本没有强制依赖apex。我的建议是不装完全不影响官方推理和评估。第二个是triton。新版mmcv在编译时如果检测到triton会尝试启用一些高版本特性算子。这些算子在老config里不一定被触发但一旦被触发容易变成“能load但推理结果全错”的隐蔽bug。保险做法是编译mmcv时不要额外安装triton或者在环境里pip uninstall triton确保mmcv跑回经典的自定义算子路径。这两个坑都是“没有报错但浪费你时间”的类型前者是编译期卡死后者是推理结果诡异写出来希望你直接避过。一路走下来我最大的感触是在4090这种新架构显卡上复现老仓库真正卡你的通常不是模型原理而是“老版本代码与新版编译链之间的错配”。拿到一个老仓库第一件事不要急着读模型代码而是先把它的依赖版本和环境声明全部列出来然后逐一对照当前显卡算力重新组合出一套可行的版本矩阵。这套思路跑通了SegFormer换到其他老仓库同样适用。
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻