告别位置编码！用SegFormer+B0/B5在Cityscapes上实战语义分割（附PyTorch代码）

告别位置编码！用SegFormer+B0/B5在Cityscapes上实战语义分割（附PyTorch代码）

当Transformer架构在计算机视觉领域掀起革命时，语义分割任务迎来了全新的可能性。传统CNN方法依赖复杂的解码器和人工设计的位置编码，而SegFormer通过创新的Hierarchical Transformer Encoder和All-MLP Decoder，不仅抛弃了位置编码的束缚，更在Cityscapes等基准数据集上实现了速度与精度的双重突破。本文将带您从零实现基于PyTorch的SegFormer-B0/B5实战，涵盖环境配置、数据预处理、训练优化到推理部署的全流程。

1. 环境配置与模型选型

1.1 硬件与软件基础要求

• GPU显存建议：
  • B0模型：最低6GB（训练阶段）
  • B5模型：最低24GB（训练阶段）
• PyTorch版本：1.9.0+（需适配CUDA 11.1）
• 关键依赖库：

  pip install mmsegmentation==0.25.0 pip install timm==0.4.12

1.2 Backbone选择策略

提示：Cityscapes数据集中，B5在1024x2048分辨率下推理速度仍可达13FPS（RTX 3090）

2. 数据预处理实战

2.1 Cityscapes数据集定制化处理

from mmseg.datasets import build_dataset from mmseg.apis import train_segmentor cfg = { 'img_dir': 'leftImg8bit/train', 'ann_dir': 'gtFine/train', 'split': 'splits/train.txt', 'type': 'CityscapesDataset', 'pipeline': [ dict(type='LoadImageFromFile'), dict(type='LoadAnnotations'), dict(type='RandomFlip', prob=0.5), dict(type='PhotoMetricDistortion'), dict(type='Normalize', mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]), dict(type='Pad', size=(512, 1024), pad_val=0), dict(type='DefaultFormatBundle'), dict(type='Collect', keys=['img', 'gt_semantic_seg']) ] }

2.2 多尺度训练技巧

• 动态缩放策略：
  • 短边随机缩放（512-1024）
  • 长边保持2048不变
• 在线数据增强：

  train_pipeline = [ dict(type='RandomMosaic', prob=0.8), dict(type='RandomRotate', degree=15), dict(type='AdjustGamma', gamma_range=(0.8, 1.2)) ]

3. 模型训练核心技巧

3.1 混合精度训练配置

# configs/segformer/segformer_mit-b0_512x1024_160k_cityscapes.py optimizer_config = dict( type='Fp16OptimizerHook', loss_scale=512., grad_clip=dict(max_norm=1, norm_type=2))

3.2 学习率调度策略

lr_config = dict( policy='poly', warmup='linear', warmup_iters=1500, warmup_ratio=1e-6, power=1.0, min_lr=0.0, by_epoch=False)

3.3 显存优化方案

• 梯度累积：每4个batch更新一次
• 激活检查点：

  model = dict( backbone=dict( use_checkpoint=True, init_cfg=dict(type='Pretrained', checkpoint='mit_b0.pth')), decode_head=dict(num_classes=19))

4. 推理部署与性能调优

4.1 测试时增强(TTA)实现

tta_model = dict( type='SegformerTTAModel', tta_cfg=dict(nms=dict(type='soft_nms', iou_thr=0.5), scales=[0.5, 0.75, 1.0]))

4.2 TensorRT加速部署

// trt_converter.py config = dict( input_shape=(1, 3, 512, 1024), precision="fp16", calib_file="calib_data.npy", dynamic_axes={ "input": {0: "batch", 2: "height", 3: "width"}, "output": {0: "batch", 2: "height", 3: "width"} })

4.3 模型量化对比

在实际道路场景测试中，B5模型处理1080p视频流时，INT8量化版本可实现35FPS的实时性能，而精度损失不到1%。这种平衡使得SegFormer非常适合自动驾驶等对实时性要求高的应用。