1. SegFormer项目概述SegFormer是NVIDIA研究院在2021年NeurIPS会议上提出的基于Transformer的语义分割模型其官方PyTorch实现已开源在GitHub。这个架构最大的特点是摒弃了传统Vision Transformer中复杂的层次化设计采用了一种新颖的混合编码器结构能够在保持轻量化的同时实现SOTA性能。我在复现这个项目时发现相比其他分割模型SegFormer在ADE20K和Cityscapes等基准数据集上展现出惊人的计算效率——B0版本仅需3.7G FLOPs就能达到78.3%的mIoU。2. 环境配置与依赖安装2.1 基础环境搭建推荐使用CUDA 11.3 PyTorch 1.12.1的组合这是经过我实测最稳定的版本配置。安装时特别注意mmcv-full的版本必须与PyTorch匹配conda create -n segformer python3.8 -y conda activate segformer pip install torch1.12.1cu113 torchvision0.13.1cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install mmcv-full1.6.1 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.12.0/index.html2.2 关键依赖版本控制SegFormer对timm库的版本极其敏感必须锁定0.3.2版本pip install timm0.3.2 pip install opencv-python4.5.1.48 git clone https://github.com/NVlabs/SegFormer.git cd SegFormer pip install -e .注意若使用更高版本的timm会导致Attention层计算异常这是源码中hardcode的shape匹配问题。3. 数据集准备与配置3.1 ADE20K数据集处理下载数据集后需要转换为MMSegmentation格式data/ade20k/ ├── images │ ├── training │ └── validation └── annotations ├── training └── validation修改config文件中的data_root路径data dict( traindict( data_rootdata/ade20k, img_dirimages/training, ann_dirannotations/training), valdict( data_rootdata/ade20k, img_dirimages/validation, ann_dirannotations/validation))3.2 自定义数据集适配若要训练自己的数据需创建继承自CustomDataset的类from mmseg.datasets.builder import DATASETS from mmseg.datasets.custom import CustomDataset DATASETS.register_module() class MyDataset(CustomDataset): CLASSES (background, object1, object2) PALETTE [[0,0,0], [128,0,0], [0,128,0]] def __init__(self, **kwargs): super().__init__( img_suffix.jpg, seg_map_suffix.png, **kwargs)4. 模型训练实战技巧4.1 多GPU训练配置使用8卡A100时的启动命令./tools/dist_train.sh \ local_configs/segformer/B2/segformer.b2.512x512.ade.160k.py \ 8 \ --work-dir work_dirs/segformer_b2 \ --load-from pretrained/mit_b2.pth关键参数调优经验batch_size建议设为每卡4-8根据显存调整lr_base6e-5时效果最佳启用SyncBN可提升2-3% mIoU4.2 学习率策略优化修改config中的lr_configlr_config dict( policypoly, warmuplinear, warmup_iters1500, warmup_ratio1e-6, power1.0, min_lr0.0, by_epochFalse)5. 模型验证与可视化5.1 多尺度测试增强./tools/dist_test.sh \ local_configs/segformer/B1/segformer.b1.512x512.ade.160k.py \ work_dirs/segformer_b1/latest.pth \ 8 \ --aug-test \ --eval mIoU5.2 预测结果可视化创建自定义调色板python demo/image_demo.py \ demo/demo.jpg \ local_configs/segformer/B1/segformer.b1.512x512.ade.160k.py \ checkpoint.pth \ --palette ade20k \ --opacity 0.56. 常见问题排查指南6.1 CUDA内存不足问题症状训练时出现CUDA out of memory 解决方案减小batch_size最低可到2启用gradient_checkpointmodel dict( backbonedict( use_checkpointTrue), decode_headdict( use_checkpointTrue))6.2 验证指标异常波动可能原因数据标注存在错误建议用labelme检查BN层未同步需设置norm_cfgdict(typeSyncBN)学习率过高尝试降低到4e-57. 模型部署优化方案7.1 ONNX导出注意事项torch.onnx.export( model, dummy_input, segformer.onnx, input_names[input], output_names[output], dynamic_axes{ input: {0: batch, 2: height, 3: width}, output: {0: batch, 1: channel}}, opset_version11)7.2 TensorRT加速技巧使用MMDeploy转换时添加优化参数python tools/deploy.py \ configs/mmseg/segmentation_tensorrt_dynamic-512x512-2048x2048.py \ segformer.b1.512x512.ade.160k.py \ checkpoint.pth \ demo/demo.jpg \ --work-dir trt_models \ --device cuda:0 \ --speed-test \ --model-optimization-options \ {fuse_bn: true, enable_fp16: true}在实际项目中我发现SegFormer-B2版本在3080Ti上使用TensorRT加速后推理速度可达45FPS512x512输入比原始PyTorch实现快3倍。对于边缘设备部署建议使用B0或B1版本并通过量化技术进一步压缩模型大小。