Win10实现Swin-Transformer 图像分割

这篇博文是关于Swin-Transformer 图像分割的应用实战，包括环境搭建、训练和测试。数据集采用ADE链接：http://data.csail.mit.edu/places/ADEchallenge/ADEChallengeData2016.zip。

Swin-Transformer 图像分割github地址：https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation

这篇文章分三个部分：

​ 第一部分介绍环境的搭建，分为Win10和Ubuntu20.04.

​ 第二部分介绍了如何配置训练参数。

​ 第三部分教大家如何配置推理参数的配置和推理结果的展示。

配置环境

win10环境配置

VS2017

pytorch 1.7.1

CUDA 11.3

1、创建虚拟环境

conda create -n swinseg python=3.7
conda activate swinseg

2、给cl.exe添加系统变量

3、安装所需要的库

pip install cython matplotlib opencv-python
pip install mmcv-full==1.3.13
pip install mmsegmentation

安装上面的库。mmcv的版本不要太高，高版本中有些参数改变了，运行的时候会有一些问题。

4、下载代码

github地址：https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation

Ubuntu20.04环境配置

Ubuntu的环境配置相对简单一些，

1、创建虚拟环境

conda create -n open-mmlab python=3.7
conda activate open-mmlab

2、安装pytorch

根据电脑的cuda版本选择pytorch，我试了1.6.0版本的可以。其他的版本在安装mmcv的时候有可能会出现问题。

3、安装mmcv-full

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

4、下载并安装Swin-Transformer-Semantic-Segmentation

git clone https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation
cd Swin-Transformer-Semantic-Segmentation 
pip install -e . #或者 pyhton setup.py develop。注意-e后面还有个. 不要丢掉。

测试环境

1、下载预训练模型

ADE20K

百度网盘的提取码是：swin

下载完后复制到项目的根目录。

2、修改./demo/image_demo.py

修改配置参数img、config、checkpoint、palette。

from argparse import ArgumentParser

from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot
from mmseg.core.evaluation import get_palette


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')
    args = parser.parse_args()

    # build the model from a config file and a checkpoint file
    model = init_segmentor(args.config, args.checkpoint, device=args.device)
    # test a single image
    result = inference_segmentor(model, args.img)
    # show the results
    show_result_pyplot(model, args.img, result, get_palette(args.palette))


if __name__ == '__main__':
    main()

修改完成后运行image_demo.py

出现上面的图说明环境没有问题了。

如果出现找不到color150.mat，百度搜索寻找一个，非常好找到。

训练

数据集配置

下载数据集然后放到./tools/data/ade、下面然后解压。

数据集路径配置在./configs/base/datasets/ade20k.py。如果不要按照我配置的路径配置，可以在这里修改路径。

模型配置

修改config/_base_/models文件夹下对应的upernet_swin.py将norm_cfg参数中的type由 SyncBN修改为BN 。将num_classes修改为150（可以不修改，我没有修改也没有出什么问题）。

修改config/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k.py的参数。
修改_base_,如下图:

_base_ = [
    '../_base_/models/upernet_swin.py', '../_base_/datasets/ade20k.py',
    '../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py'
]

默认是ade20k，如果选用其他的数据集，则修改对应的py脚本、比如pascal_voc12数据集

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

修改所有num_classes，ade的类别是150。

修改data['samples_per_gpu']，这个就是batchsize，不能小于2。

修改train.py

通过from mmseg import __version__这句话找到mmesg适用的版本，然后将其MAX修改为1.3.13。不然会有版本的问题

参照下面的配置参数修改：

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')
    group_gpus = parser.add_mutually_exclusive_group()
    group_gpus.add_argument(
        '--gpus',
        type=int,
        help='number of gpus to use '
        '(only applicable to non-distributed training)')
    group_gpus.add_argument(
        '--gpu-ids',
        type=int,
        nargs='+',
        help='ids of gpus to use '
        '(only applicable to non-distributed training)')
    parser.add_argument('--seed', type=int, default=None, help='random seed')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')
    parser.add_argument(
        '--options', nargs='+', action=DictAction, help='custom options')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args

然后运行train.py

测试

修改./demo/image_demo.py 中的checkpoint路径,然后运行即可。

参考：

Win10配置Swin-Transformer-Semantic-Segmentation并训练自己数据集_哔哩哔哩_bilibili