Swin-Transformer训练自己的数据集

前言
一、虚拟环境搭建
二、制作自己的数据集
- 1.样本准备
- 2.调试代码
总结

前言

Swin-Transformer精度较高，但其对显卡要求同样较高，我的是RTX2070，8G显存，当设置crop size为512512时，batchsize设置为2，才不会报OOM。当crop size为10241024时，如果类别较少，还可以跑通，类别较多则会报错。

一、虚拟环境搭建

1.下载代码

首先下载Swin-Transformer的源码，在PyCharm中创建针对Swin-Transformer算法的虚拟环境。此处不再多说。

2.相关库安装

在PyCharm中创建好虚拟环境后，并打开Terminal，会自动进入虚拟环境，之后执行命令会将库安装在虚拟环境中。安装requirement.txt中所需的库，自动安装的话有时候版本会出现问题。如PyTorch需要1.6.0，且需要和自己电脑的CUDA版本对应，所以自行安装如下

pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

安装mmcv-full 1.1.5版本，也可参照mmcv安装文档安装

pip install mmcv-full==1.1.5 -f https://download.openmmlab.com/mmcv/dist/cu101/torch1.6.0/index.html

安装mmseg

pip install mmsegmentation

安装代码包中requirements.txt所包含的库，

pip install -r requirements\docs.txt

打开py文件后，PyCharm也会自动提示按照requirement安装，按照该提示安装亦可。

3.环境测试

安装完成后，为了测试环境是否符合要求，先用给定demo对图片进行预测，能跑通则说明环境大体可用。
首先在代码地址下载预训练模型，本文从百度云下载ADE20K数据集Swin-T对应的预训练模型，密码swin。亦可从代码地址下载其他模型。将下载的模型放在demo文件夹中。
打开demo目录下image_demo.py，修改代码如下：

def main():
    parser = ArgumentParser()
    parser.add_argument('--img', default='demo.png', help='Image file')
    parser.add_argument('--config', default='../configs/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k.py',help='Config file')
    parser.add_argument('--checkpoint', default='upernet_swin_tiny_patch4_window7_512x512.pth',help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--palette',
        default='ade20k',
        help='Color palette used for segmentation map')

然后运行demo.py，如果可以弹出预测结果如下，则说明环境搭建基本成功。
在这里插入图片描述

二、制作自己的数据集

1.样本准备

Swin-Transformer支持VOC格式数据集，本人用Labelme标注的样本如下，转为VOC格式即可。
在这里插入图片描述
在labelme中执行如下命令，其中LabelmeDataPath 表示labelme标注好的样本所在文件夹；VOCDataPath 表示输出的voc格式的文件夹，注意该文件夹不能手动创建，如果该文件夹已存在则提示已存在；labels.txt中保存的是所有样本类别名称，一个类别一行
labelme2voc.py --input_dir LabelmeDataPath --output_dir VOCDataPath --labels labels.txt

将生成的VOC格式数据复制到如下图所示文件夹下，不存在的目录可自行手动创建。在这里插入图片描述
JPEGImages中存放所有图片，如下所示：

SegmentationClass文件夹中是VOC格式的样本结果图

创建如下目录，train.txt中是训练用的图片名称，一行一个图片；val.txt是验证所需图片的名称。在这里插入图片描述

2.调试代码

在mmseg/datasets/voc.py中根据自己样本情况修改类别和颜色，类别和颜色数要一致。

代码如下（示例）：

    CLASSES = ('background', 'cat', 'dag')

    PALETTE = [[120, 120, 120], [6, 230, 230], [128,192,0]]

在configs/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k.py中修改类别数，搜索num_classes，改为自己的类别数，这里修改为3。此外该算法的batch_size 也是在此文件中修改，最后一行

    data = dict(sample_per_gpu = 2)

此处的2即为batch_size，我的显存为8G，cropsize为512，可以设置为2，当cropsize为1024时，设置为1，但当我的类别增加很多时，batchsize为1显存也爆了。所以根据自身情况修改即可。
由于本文用VOC格式样本，所以将此处ade20k.py修改为voc格式对应的文件。此外，由于只是测试，不需要迭代太多次，将schdules_160k.py修改为schdules_20k.py，如下

    _base_ = [
     '../_base_/models/upernet_swin.py', '../_base_/datasets/pascal_voc12.py',
     '../_base_/default_runtime.py', '../_base_/schedules/schedule_2k.py'
 	]

由于schduls_20k.py还是需要迭代2万次，如还想减少最大迭代次数，可在改文件中进行修改最大迭代次数，如下所示，注意：下边的3000，表示没迭代3000次会输出一次模型文件，如最大迭代次数小于该数，输出目录下只会有日志文件，而没有权重模型。在这里插入图片描述
在configs_base_\datasets\pascal_voc12.py中，crop_size可用于修改输入图片的大小，默认为512。
如果你是单GPU训练，在configs_base_\models\upernet_swin.py中将type='SyncBN’修改为type=‘BN’，否则可能回报错。另外此处的num_classes同样需要修改为自身样本的类别数，本文为3
修改tools\train.py

    def parse_args():
    parser = argparse.ArgumentParser(description='Train a segmentor')
    parser.add_argument('--config',default='../configs/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k.py', help='train config file path')
    parser.add_argument('--work-dir',default='output', help='the dir to save logs and models')
    parser.add_argument(
        '--load-from',default='../upernet_swin_tiny_patch4_window7_512x512.pth', help='the checkpoint file to load weights from')
    parser.add_argument(
        '--resume-from', help='the checkpoint file to resume from')
    parser.add_argument(
        '--no-validate',
        action='store_true',
        help='whether not to evaluate the checkpoint during training')
 	]

运行tools\train.py，训练完成后，会在tools文件夹下生成模型文件，可将生成的模型文件用demo.py测试训练结果

总结

以上是Swin-Transformer的训练记录。