前言
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测试训练结果文章来源：https://www.toymoban.com/news/detail-446184.html
总结
以上是Swin-Transformer的训练记录。文章来源地址https://www.toymoban.com/news/detail-446184.html
到了这里，关于Swin-Transformer训练自己的数据集的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
