ICode9

精准搜索请尝试: 精确搜索
首页 > 其他分享> 文章详细

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

2021-07-05 11:58:50  阅读:1076  来源: 互联网

标签:transformer Swin 数据 py 修改 github 图像 data


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

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

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

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

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

data     /    class_1

                 class_2

                 class_3

                    .......

3. 修改 main.py中的参数,主要修改--cfg 、data_path 、batch-size、resume

resume就是载入官方的预训练权重,给它路径即可,代码中我已修改最终输出的类别通道数。其他的自由修改,建议默认。

4. 打开config.py文件,修改

_C.SAVE_FREQ = 10 ,每多少个epoch保存一次模型
_C.TRAIN.EPOCHS = 300  总共训练多少个epoch

5. 修改 data/build.py中第69行,num_classes=..自己的类别数

6. 运行命令 python -m torch.distributed.launch  --nproc_per_node 1(gpu数量) --master_port 12345 main.py 即可

标签:transformer,Swin,数据,py,修改,github,图像,data
来源: https://blog.csdn.net/hi_gril/article/details/118486070

本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享;
2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关;
3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关;
4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除;
5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

专注分享技术,共同学习,共同进步。侵权联系[81616952@qq.com]

Copyright (C)ICode9.com, All Rights Reserved.

ICode9版权所有