从零开始使用MMSegmentation训练Segformer

从零开始使用MMSegmentation训练Segformer

写在前面：最新想要用最新的分割算法如：Segformer or SegNeXt 在自己的数据集上进行训练，但是有不是搞语义分割出身的，而且也没有系统的学过MMCV以及MMSegmentation。所以就折腾了很久，感觉利用MMSegmentation搭建框架可能比较系统，但是对于不熟悉的或者初学者非常不友好，因此记录一下自己training Segformer的心路历程。

Segformer paper: https://arxiv.org/abs/2105.15203>
官方实现: https://github.com/NVlabs/SegFormer>
纯Torch版Segformer: https://github.com/camlaedtke/segmentation_pytorch>

方法

由于本人不是研究语义分割的，所以只能简要地介绍一下Segformer。
SegFormer的动机在于：

① ViT作为backbone只能输出固定分辨率的特征图，这对于密集预测任务显然不够友好；
② 由于self-attention操作的存在，transformer的运算量和参数两都非常大，不利于大尺度图像的分割。

为此作者提出了相应的创新：

① 先是对transformer进行层次化结构设计，得到多层级的特征图；
② 构造轻量级的decoder，仅使用MLP进行特征聚合。
③ 除此之外，SegFormer抛弃了位置信息编码，选择采用MixFCN来学习位置信息，这样可以很好地扩充到不同尺度的测试环境下（避免由于尺寸变化，需要对positional-encoding进行插值，从而影响性能）。最后提出的模型在ADE20k上达到了新sota，并且在速度、性能和鲁棒性上都表现很好。

程序复现

在重新训练过程中主要参考了：手把手教你使用Segformer训练自己的数据
作者给的教程比较详细， 但是有几处修改并不合适，导致我复现出来的结果没啥效果，因此记录一下自己的采坑记录。
自己的主要配置为：

CUDA 10.1
Pytorch 1.10.0, torchvision 0.11.1
MMCV-full 1.3.0

其中在安装MMCV-full过程中还遇到了很多问题，主要是版本不适配的原因导致的。
在安装好环境后，首先从Github下载SegFormer的项目工程： https://github.com/NVlabs/SegFormer
然后进去SegFormer目录：

pip install -r requirements.txt
pip install -e . --use

安装需要的依赖。

数据集准备

代码默认用的是ADE20K数据集进行训练
ADE20K数据集 格式如下，按照要求放就完了

├── data
│   ├── ade
│   │   ├── ADEChallengeData2016
│   │   │   ├── annotations
│   │   │   │   ├── training
│   │   │   │   ├── validation
│   │   │   ├── images
│   │   │   │   ├── training
│   │   │   │   ├── validation

但是@中科哥哥使用的是VOC的数据格式，因此就使用了VOC的数据格式

├── VOCdevkit
│   ├── VOC2012
│   │   ├── ImageSets
│   │   │   ├── Segmentation
│   │   │   │   ├── train.txt
│   │   │   │   ├── val.txt
│   │   │   │   ├── trainval.txt│   │   │   ├── JPEGImages
│   │   │   │   ├── *.jpg    #所有图片│   │   │   ├── SegmentationClass
│   │   │   │   ├── *.jpg    #所有标签图

在这里可以根据自己的需要修改
下面是我自己的数据格式：

├── VOCdevkit
│   ├── VOC2012
│   │   ├── ImageSets
│   │   │   ├── Segmentation
│   │   │   │   ├── train.txt
│   │   │   │   ├── val.txt
│   │   │   │   ├── test.txt│   │   │   ├── JPEGImages
│   │   │   │   ├── *.png#所有图片│   │   │   ├── SegmentationClass
│   │   │   │   ├── *.png    #所有标签图

其实完全可以简洁一点：

├── MFNet
│   ├── Segmentation
│   │   ├── train.txt
│   │   ├── val.txt
│   │   ├── test.txt│   ├── Images
│   │   ├── *.png#所有图片│   ├── Label
│   │   ├── *.png    #所有标签图

其中： train.txt； val.txt； test.txt； 只要图片名，不需要后缀和路径 如下

后面的程序修改都基于复杂的版本进行介绍（自己在程序复现时使用的复杂的目录，因为是按照的教程来的）
数据准备好之后可以在SegFormer目录先新建一个/datasets/ 目录来存放自己的数据集

程序修改

1. 在 mmseg/datasets/voc.py修改自己数据集的类别即修改CLASSES 和 PALETTE在我自己的数据集中一共由于9个类别，所以修改如下：
   

2. 在 mmseg/models/decode_heads/segformer_head.py 中BatchNorm 方式（如果使用单卡训练的话就修改，多卡训练的话就不用修改）。 将第59行的SyncBN 修改为 BN
   

3. 修改 local_configs/segformer/B5/segformer.b5.640x640.ade.160k.py 的配置文件(这里我们使用的是B5模型，需要使用哪个模型就修改对应的配置文件即可，配置文件都位于：**local_configs/segformer/**下 )；主要修改

__base__=[]中的数据集文件路径（也就是下图中的第二行）
指定dataset_type的类型，此处
dataset_type = 'PascalVOCDataset'
data_root = '/data1/timer/Segmentation/SegFormer/datasets/VOC2012'  也可以给相对路径。
然后根据自己的数据需要修改文件中的crop_size, train_pipline中的img_scale，以及test_pipline中的img_scale

同时好需要在data字典中指定 img_dir, ann_dir, 以及split的路径，如果是单卡训练的话需要将norm_cfg 的type由的SyncBN 修改为 BN

接下来继续修改模型相关的文件，主要是给定预训练权重的位置即修改：pretrained 以及backbone[‘type’]，这里的type因为使用的是B5的结构所以type就指定为mit_b5，然后预训练权重需要从项目中给定的链接下载。值得注意的是还需要指定decode_head[‘num_classes’] (这个需要根据你的数据集来指定，因为我的数据集中包含9类，所以这里就设置为9了)

4. 在 local_configs/base/models/segformer.py 修改

norm_cfg[‘type’]=‘BN’
num_classes=9 (这里修改成你数据集对应的类别的数量)

5. 在 tools/train.py中修改

parser.add_argument('--config', default='/data1/timer/Segmentation/SegFormer/local_configs/segformer/B5/segformer.b5.640x640.ade.160k.py') 
parser.add_argument('--work-dir', default='res_MFNet')

其中 /data1/timer/Segmentation/SegFormer/local_configs/segformer/B5/segformer.b5.640x640.ade.160k.py 是配置文件的路径
res_MFNet是训练日志和模型保存的路径
同时指定GPU的卡号

    group_gpus.add_argument('--gpu-ids',type=int,        default=[0],help='ids of gpus to use ''(only applicable to non-distributed training)')

6. 进入tools目录下运行

python train.py

即可开始训练模型。
由于本人也在摸索阶段，有不当之处，恳请各位不吝赐教。也欢迎大家交流：[e-mail]