我们希望使用swin-transformer算法实现物体的分类。
swin-transformer的github地址为:https://github.com/microsoft/Swin-Transformer
本文参考:Swin Transformer实战:使用 Swin Transformer实现图像分类-阿里云开发者社区
(1)配置要求
windows版本:windows11
pytorch版本:1.7.1
python版本:3.7.3(至少要大于3.6.2,因为pytorch1.7.1的python最低要求是3.6.2)
cuda版本:11.0(pytorch1.7.1在windows11下使用,最少需要cuda11.0)
以上配置为我试验swin-transformer运行的相对比较低的配置要求。
(2)安装方式
torch1.7.1的安装命令:pip install torch==1.7.1 -f https://download.pytorch.org/whl/torch_stable.html
torchvision的安装命令:pip install torchvision==0.8.2 -f https://download.pytorch.org/whl/torch_stable.html
cuda中官网下载toolkit即可,windows下cuda11.0.X版本可选择windows10对应的版本。
swin-transformer默认读取imagenet格式的数据集。
数据集的目录结构如下:
1、修改config.py文件
_C.DATA.DATA_PATH = 'D:\\temp\\pic_ai\\swin_transformer_test'
_C.MODEL.NUM_CLASSES = 2
_C.DATA.NUM_WORKERS = 0
_C.DATA.PIN_MEMORY = False
2、修改build.py文件
将nb_classes =1000改为nb_classes = config.MODEL.NUM_CLASSES,如下所示:
将部分_pil_interp修改为str_to_pil_interp,如下图所示:
3、修改utils.py文件
由于类别默认是1000,所以加载模型的时候会出现类别对不上的问题,所以需要修改load_checkpoint方法。在加载预训练模型之前增加修改预训练模型的方法:
if checkpoint['model']['head.weight'].shape[0] == 1000:
checkpoint['model']['head.weight'] = torch.nn.Parameter(
torch.nn.init.xavier_uniform(torch.empty(config.MODEL.NUM_CLASSES, 768)))
checkpoint['model']['head.bias'] = torch.nn.Parameter(torch.randn(config.MODELNUM_CLASSES))
msg = model.load_state_dict(checkpoint['model'], strict=False)
4、修改main.py文件
(1)将如下代码注释:
(2)将torch.distributed.init_process_group修改为:
torch.distributed.init_process_group('gloo', init_method='file://tmp/somefile', rank=0, world_size=1)
该函数只有在pytorch1.7.1以上才支持。
5、修改lr_scheduler.py文件
将如下代码注释掉
五、运行训练命令:
python.exe D:/workspace/transformer/Swin-Transformer/main.py --cfg configs/swin/swin_tiny_patch4_window7_224.yaml --local_rank 0 --batch-size 2
执行后显示如下: