SegFormer实战:从零部署到ADE20K语义分割

SegFormer实战:从零部署到ADE20K语义分割 1. 环境搭建从零配置SegFormer开发环境第一次接触SegFormer时最头疼的就是环境配置。记得去年我在Ubuntu 18.04上折腾了两天才跑通第一个demo现在把这些经验都总结成最简步骤。SegFormer基于PyTorch和MMSegmentation框架对环境版本要求比较严格建议完全按照以下步骤操作。先创建conda环境这是避免依赖冲突的最佳实践conda create -n seg python3.8 -y conda activate seg安装PyTorch时要注意CUDA版本匹配。如果你的显卡是RTX 30系列需要CUDA 11pip install torch1.7.0cu110 torchvision0.8.1cu110 torchaudio0.7.0 -f https://download.pytorch.org/whl/torch_stable.htmlMMCV是MMSegmentation的核心依赖必须安装完整版pip install mmcv-full1.3.0 -f https://download.openmmlab.com/mmcv/dist/cu110/torch1.7.0/index.html其他必要依赖包括OpenCV和交互式Python环境pip install opencv-python4.5.1.48 IPython timm最后安装SegFormer本体。这里有个小技巧先用清华源加速安装基础依赖再用开发模式安装项目cd SegFormer pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple pip install -e . --user验证安装是否成功可以运行import mmseg print(mmseg.__version__) # 应该输出1.3.0常见踩坑点如果遇到GLIBCXX_3.4.26 not found错误需要更新libstdc6Windows用户建议使用WSL2原生Windows环境会有各种路径问题3090等新显卡需要安装torch 1.8和CUDA 11.12. 数据准备ADE20K数据集处理全攻略ADE20K是MIT发布的语义分割基准数据集包含2万张图像和150类标注。原始数据需要特殊处理才能被SegFormer使用这个过程我踩过不少坑。首先下载和解压数据集wget http://data.csail.mit.edu/places/ADEchallenge/ADEChallengeData2016.zip unzip ADEChallengeData2016.zip关键步骤是调整目录结构。SegFormer要求特定路径格式很多人在这里出错data └── ade └── ADEChallengeData2016 ├── annotations │ ├── training │ └── validation └── images ├── training └── validation用以下命令快速完成结构调整mkdir -p data/ade mv ADEChallengeData2016 data/ade/配置文件通常在SegFormer/configs/_base_/datasets/ade20k.py需要检查以下关键参数dataset_type ADE20KDataset data_root data/ade/ADEChallengeData2016 data dict( samples_per_gpu4, # 根据GPU显存调整 workers_per_gpu4, # 根据CPU核心数调整 traindict( typedataset_type, data_rootdata_root, img_dirimages/training, ann_dirannotations/training ), # ...其他配置 )数据增强策略对分割效果影响很大。建议在pipeline中添加这些变换train_pipeline [ dict(typeLoadImageFromFile), dict(typeLoadAnnotations), dict(typeRandomFlip, prob0.5), dict(typePhotoMetricDistortion), # 亮度/对比度扰动 dict(typeNormalize, mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375]), dict(typeDefaultFormatBundle), dict(typeCollect, keys[img, gt_semantic_seg]) ]3. 模型训练从预训练权重到自定义训练SegFormer有B0-B5六个版本B1在精度和速度上比较平衡。训练过程需要注意以下关键点。首先下载ImageNet预训练权重mkdir pretrained cd pretrained wget https://huggingface.co/nvidia/mit-b1/resolve/main/mit_b1.pth修改模型配置有两个重点将SyncBN改为普通BN单卡训练时调整解码头参数在segformer.b1.512x512.ade.160k.py中修改norm_cfg dict(typeBN, requires_gradTrue) # 修改SyncBN为BN model dict( pretrainedpretrained/mit_b1.pth, backbonedict(typemit_b1), decode_headdict( num_classes150, # ADE20K有150类 channels256, # 增大特征通道数 loss_decodedict( typeCrossEntropyLoss, use_sigmoidFalse, loss_weight1.0 ) ) )启动训练前解决常见报错# yapf版本冲突 pip install yapf0.40.1 # 指定训练GPU sed -i s/defaultNone/default[0]/g tools/train.py开始训练的命令如下建议使用nohup保持会话python tools/train.py \ local_configs/segformer/B1/segformer.b1.512x512.ade.160k.py \ --work-dir work_dirs/exp1 \ --deterministic \ --seed 42训练过程监控技巧用TensorBoard查看损失曲线tensorboard --logdir work_dirs当验证mIoU不再提升时可以提前终止8GB显存建议设置batch_size416GB可以设为84. 模型调优提升分割精度的实用技巧经过多次实验我总结了这些提升SegFormer效果的实用方法。学习率策略对结果影响显著。推荐使用多项式衰减optimizer dict( typeAdamW, lr6e-5, betas(0.9, 0.999), weight_decay0.01 ) lr_config dict( policypoly, power1.0, min_lr0.0, by_epochFalse )数据增强的组合很关键这套配置在我测试中效果最好train_pipeline [ dict(typeRandomFlip, flip_ratio0.5), dict(typeRandomRotate, degree(-20, 20)), dict(typeAlbu, transforms[ dict(typeRandomBrightnessContrast, brightness_limit0.3, contrast_limit0.3), dict(typeRGBShift, r_shift_limit15, g_shift_limit15, b_shift_limit15) ]), dict(typeNormalize, mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375]) ]针对小目标的改进方案使用更大的输入分辨率如640x640在decode_head中添加辅助损失decode_headdict( typeSegFormerHead, loss_auxdict( typeCrossEntropyLoss, loss_weight0.4 ) )混合精度训练能大幅节省显存export AMPtrue python tools/train.py ...模型量化部署方案model dict( backbonedict( typemit_b1, output_attentionFalse # 推理时不需要attention map ), test_cfgdict( modeslide, crop_size(512, 512), stride(256, 256) ) )5. 模型评估与可视化验证你的成果训练完成后我们需要科学评估模型性能。ADE20K的标准指标是mIoU平均交并比。首先下载验证集权重mkdir checkpoints cd checkpoints wget https://huggingface.co/nvidia/segformer-b1-finetuned-ade-512-512/resolve/main/pytorch_model.bin mv pytorch_model.bin segformer.b1.512x512.ade.160k.pth修改评估代码避免数值溢出# 在mmseg/core/evaluation/metrics.py中 total_area_intersect np.zeros((num_classes, ), dtypenp.float64) total_area_union np.zeros((num_classes, ), dtypenp.float64)运行验证命令python tools/test.py \ local_configs/segformer/B1/segformer.b1.512x512.ade.160k.py \ checkpoints/segformer.b1.512x512.ade.160k.pth \ --eval mIoU \ --show-dir results可视化结果时推荐使用Cityscapes调色板from mmseg.apis import inference_segmentor, init_segmentor import mmcv model init_segmentor(config_file, checkpoint_file, devicecuda:0) img mmcv.imread(demo.png) result inference_segmentor(model, img) # 自定义可视化 palette [[120, 120, 120], [180, 120, 120], ...] # 150类颜色 draw_img model.show_result(img, result, palettepalette, opacity0.5) mmcv.imwrite(draw_img, result.jpg)性能优化技巧启用cudnn benchmark加速推理torch.backends.cudnn.benchmark True使用TensorRT部署能获得3-5倍加速对于视频流处理可以复用中间特征减少计算量