官方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环境，根据官方github操作即可，当然在安装apex时可能会遇到坑，基本是torch.verson.cuda版本与本机安装的cuda版本不一致导致，可调整torch的版本适应本机的cuda版本。
2 . 准备数据集，代码使用的是ImageFolder函数生成dataloader，因此，你的数据集应该有如下格式：
data / class_1
class_2
class_3
…….
（class_1，class_2 。。。是类别文件夹的名字。）
- 修改 main.py中的参数，主要修改–cfg 、data_path 、batch-size、resume
resume就是载入官方的预训练权重，给它路径即可，代码中我已修改最终输出的类别通道数，需要自己调整，在main.py的
其他的自由修改，建议默认。
- 打开config.py文件，修改
_C.SAVE_FREQ = 10 ，每多少个epoch保存一次模型
_C.TRAIN.EPOCHS = 300 总共训练多少个epoch
- 修改 data/build.py中第69行，num_classes=..自己的类别数，与3中修改的main.py保持类别一致。
- 运行命令 python -m torch.distributed.launch –nproc_per_node 1 （gpu数量） –master_port 12345 main.py 即可
Original: https://blog.csdn.net/hi_gril/article/details/118486070
Author: 这颗
Title: Swin-transformer训练自己的数据集，图像分类
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/688666/
转载文章受原作者版权保护。转载请注明原作者出处！