Swin-transformer训练自己的数据集,图像分类

人工智能159

官方swin-transformer目前只支持训练imageNet,其次加载预训练有点东西,想要训练自己的数据集需要稍微改改代码,为了方便使用,我将改好的上传到github,大家仅需要自己修改参数即可使用。github: https://github.com/sunanlin13174/Image-train-Swin-transformer

我的数据集是图案比较丰富的,但是部分类或者数据集本身数据不多,因此在官方已有数据增强:亮度、对比度、饱和度变化、mixup、裁剪等基础上,添加了上下左右中心裁剪并翻转,生成更多的数据,然后以0.5的概率将其转为灰度图(r=g=b),这部分代码可在data/build.py中修改,后期使用flask部署到服务器端运行。

Swin-transformer训练自己的数据集,图像分类

其次,主要讲一下如何使用这份代码。

  1. 适配swin-transformer环境,根据官方github操作即可,当然在安装apex时可能会遇到坑,基本是torch.verson.cuda版本与本机安装的cuda版本不一致导致,可调整torch的版本适应本机的cuda版本。

2 . 准备数据集,代码使用的是ImageFolder函数生成dataloader,因此,你的数据集应该有如下格式:

data / class_1

class_2

class_3

.......

输入验证码查看隐藏内容

扫描二维码关注本站微信公众号 Johngo学长
或者在微信里搜索 Johngo学长
回复 svip 获取验证码
wechat Johngo学长