前言
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模板网!