从零开始使用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 . --user
安装需要的依赖。
数据集准备
代码默认用的是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/ 目录来存放自己的数据集
程序修改
-
在 mmseg/datasets/voc.py修改自己数据集的类别即修改CLASSES 和 PALETTE在我自己的数据集中一共由于9个类别,所以修改如下:
-
在 mmseg/models/decode_heads/segformer_head.py 中BatchNorm 方式(如果使用单卡训练的话就修改,多卡训练的话就不用修改)。 将第59行的SyncBN 修改为 BN
-
修改 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了)
- 在 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目录下运行文章来源:https://www.toymoban.com/news/detail-484590.html
python train.py
即可开始训练模型。
由于本人也在摸索阶段,有不当之处,恳请各位不吝赐教。也欢迎大家交流:2458707789@qq.com文章来源地址https://www.toymoban.com/news/detail-484590.html
到了这里,关于从零开始使用MMSegmentation训练Segformer的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!