Swin Transformer簡介

目標(biāo)檢測刷到58.7 AP！

實例分割刷到51.1 Mask AP！

語義分割在ADE20K上刷到53.5 mIoU！

今年，微軟亞洲研究院的Swin Transformer又開啟了吊打CNN的模式，在速度和精度上都有很大的提高。這篇文章帶你實現(xiàn)Swin Transformer圖像分類。

資料匯總

論文: https://arxiv.org/abs/2103.14030

代碼: https://github.com/microsoft/Swin-Transformer

論文翻譯：https://wanghao.blog.csdn.net/article/details/120724040

一些大佬的B站視頻：

1、霹靂吧啦Wz：https://www.bilibili.com/video/BV1yg411K7Yc?from=search&seid=18074716460851088132&spm_id_from=333.337.0.0

2、ClimbingVision社區(qū)：震驚！這個關(guān)于Swin Transformer的論文分享講得太透徹了！_嗶哩嗶哩_bilibili

關(guān)于Swin Transformer的資料有很多，在這里就不一一列舉了，我覺得理解這個模型的最好方式：源碼+論文。

環(huán)境配置

1、電腦環(huán)境：

操作系統(tǒng)：win10

CUDA版本：11.2

2、創(chuàng)建虛擬環(huán)境swin

conda create -n swin python=3.7 -y conda activate swin

3、安裝pytorch

conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=11.0 -c pytorch

4、安裝timm

pip install timm==0.3.2

5、安裝apex

APEX是英偉達開源的，完美支持PyTorch框架，用于改變數(shù)據(jù)格式來減小模型顯存占用的工具。其中最有價值的是amp（Automatic Mixed Precision），將模型的大部分操作都用Float16數(shù)據(jù)類型測試，一些特別操作仍然使用Float32。并且用戶僅僅通過三行代碼即可完美將自己的訓(xùn)練代碼遷移到該模型。實驗證明，使用Float16作為大部分操作的數(shù)據(jù)類型，并沒有降低參數(shù)，在一些實驗中，反而由于可以增大Batch size，帶來精度上的提升，以及訓(xùn)練速度上的提升。

5.1 下載apex

網(wǎng)址 https://github.com/NVIDIA/apex，下載到本地文件夾。解壓后進入到apex的目錄安裝依賴。在執(zhí)行命令；

cd C:\Users\WH\Downloads\apex-master #進入apex目錄
pip install -r requirements.txt

5.2 安裝apex

依賴安裝完后，打開cmd，cd進入到剛剛下載完的apex-master路徑下，運行：

python setup.py install

然后跑了一堆東西，最后是這樣的：

安裝完成！

6、安裝一些其他的包

pip install opencv-python==4.4.0.46 termcolor==1.1.0 yacs==0.1.8

數(shù)據(jù)集

數(shù)據(jù)集采用最經(jīng)典的貓狗大戰(zhàn)數(shù)據(jù)集。數(shù)據(jù)集地址：鏈接：https://pan.baidu.com/s/1ZM8vDWEzgscJMnBrZfvQGw 提取碼：48c3

如果連接失效請聯(lián)系我，或者你也可以從別的途徑獲得。

項目結(jié)構(gòu)

使用tree命令打印整個項目的結(jié)構(gòu)

Swin-Transformer-main
    ├─configs#配置文件 ├─data#處理數(shù)據(jù)集相關(guān)的操作 │ 
    ├─dataset #數(shù)據(jù)集結(jié)構(gòu) │  ├─test │  ├─train
    │  │  ├─cat
    │  │  └─dog
    │  └─val
    │      ├─cat
    │      └─dog
    ├─figures
    ├─models#Swin的模型文件 │ 
    ├─output#訓(xùn)練模型的輸出

訓(xùn)練

1、獲取代碼和預(yù)訓(xùn)練模型

從https://github.com/microsoft/Swin-Transformer下載代碼，然后放到本地。然后解壓。

在get_started.md找到預(yù)訓(xùn)練模型下載路徑，下載下來然后放到Swin-Transformer根目錄。

2、制作數(shù)據(jù)集

構(gòu)建數(shù)據(jù)集，數(shù)據(jù)集結(jié)構(gòu)如下：

dataset #數(shù)據(jù)集結(jié)構(gòu)
    ├─test ├─train
    │  ├─cat │  └─dog
    └─val
           ├─cat └─dog

從原數(shù)據(jù)集中取出一部分數(shù)據(jù)集放入train對應(yīng)的類別中，一部分放入val對應(yīng)的類別中。把原數(shù)據(jù)集中的test直接復(fù)制到test中。

3、修改config.py文件

_C.DATA.DATA_PATH = 'dataset' # Dataset name _C.DATA.DATASET = 'imagenet' # Model name _C.MODEL.NAME = 'swin_tiny_patch4_window7_224' # Checkpoint to resume, could be overwritten by command line argument
_C.MODEL.RESUME ='swin_tiny_patch4_window7_224.pth' # Number of classes, overwritten in data preparation
_C.MODEL.NUM_CLASSES = 2

對上面參數(shù)的解釋：

_C.DATA.DATA_PATH ：數(shù)據(jù)集路徑的根目錄，我定義為dataset。

_C.DATA.DATASET：數(shù)據(jù)集的類型，這里只有一種類型imagenet。

_C.MODEL.NAME：模型的名字，對應(yīng)configs下面yaml的名字，會在模型輸出的root目錄創(chuàng)建對應(yīng)MODEL.NAME的目錄。

_C.MODEL.RESUME：預(yù)訓(xùn)練模型的目錄。

_C.MODEL.NUM_CLASSES：模型的類別，默認是1000，按照數(shù)據(jù)集的類別數(shù)量修改。

4、修改build.py

將nb_classes =1000改為nb_classes = config.MODEL.NUM_CLASSES

5、修改utils.py

由于類別默認是1000，所以加載模型的時候會出現(xiàn)類別對不上的問題，所以需要修改load_checkpoint方法。在加載預(yù)訓(xùn)練模型之前增加修改預(yù)訓(xùn)練模型的方法：

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)

6、修改main.py

將92-94注釋，如下圖：

將312行修改為：torch.distributed.init_process_group('gloo', init_method='file://tmp/somefile', rank=0, world_size=1)

7、運行訓(xùn)練命令

打開Terminal，運行如下命令：

python main.py --cfg configs/swin_tiny_patch4_window7_224.yaml --local_rank 0 --batch-size 16

如果想單獨驗證，運行命令：

python  main.py --eval --cfg configs/swin_tiny_patch4_window7_224.yaml --resume ./output/swin_tiny_patch4_window7_224/default/ckpt_epoch_1.pth --data-path dataset --local_rank 0

推理

這個項目沒有推理腳本，我自己寫了一個。寫這部分需要看懂驗證部分的代碼即可。

1、導(dǎo)入包和配置參數(shù)

import torch.utils.data.distributed
import torchvision.transforms as transforms from PIL import Image from torch.autograd import Variable
import os from models import build_model from config import get_config
import argparse

def parse_option():
    parser = argparse.ArgumentParser('Swin Transformer Test script', add_help=False)
    parser.add_argument('--cfg', default='configs/swin_tiny_patch4_window7_224.yaml', type=str, metavar="FILE", help='path to config file', )
    parser.add_argument( "--opts", help="Modify config options by adding 'KEY VALUE' pairs. ", default=None, nargs='+',
    )

    # easy config modification
    parser.add_argument('--batch-size', type=int, help="batch size for single GPU")
    parser.add_argument('--data-path', type=str, help='path to dataset')
    parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
    parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'], help='no: no cache, ' 'full: cache all data, ' 'part: sharding the dataset into nonoverlapping pieces and only cache one piece')
    parser.add_argument('--resume', default='output/swin_tiny_patch4_window7_224/default/ckpt_epoch_1.pth', help='resume from checkpoint')
    parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps")
    parser.add_argument('--use-checkpoint', action='store_true', help="whether to use gradient checkpointing to save memory")
    parser.add_argument('--amp-opt-level', type=str, default='O1', choices=['O0', 'O1', 'O2'], help='mixed precision opt level, if O0, no amp is used')
    parser.add_argument('--output', default='output', type=str, metavar='PATH', help='root of output folder, the full path is <output>/<model_name>/<tag> (default: output)')
    parser.add_argument('--tag', help='tag of experiment')
    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
    parser.add_argument('--throughput', action='store_true', help='Test throughput only')
    parser.add_argument("--local_rank", default='0', type=int, help='local rank for DistributedDataParallel')
    args, unparsed = parser.parse_known_args()  config = get_config(args)

    return args, config 

這個配置參數(shù)是為了創(chuàng)建模型，從main.py中復(fù)制過來，然后將required=True這樣的字段刪除。

定義class、創(chuàng)建transform

transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
classes = ("cat", "dog")

將圖像resize為224×224大小

定義類別，順序和數(shù)據(jù)集對應(yīng)。

2、創(chuàng)建模型

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
_, config = parse_option()
model = build_model(config) checkpoint = torch.load('output/swin_tiny_patch4_window7_224/default/ckpt_epoch_1.pth', map_location='cpu')
model.load_state_dict(checkpoint['model'], strict=False)
model.eval()
model.to(DEVICE)

判斷gpu是否可用，如果不可以使用cpu。

獲取config參數(shù)

創(chuàng)建模型

加載訓(xùn)練的模型權(quán)重

將權(quán)重放入model中。

3、開始推理

定義測試集的路徑，然后循環(huán)預(yù)測每張圖片

path = 'dataset/test/'
testList = os.listdir(path) for file in testList:
    img = Image.open(path + file)
    img = transform_test(img)
    img.unsqueeze_(0)
    img = Variable(img).to(DEVICE) out = model(img)
    # Predict _, pred = torch.max(out.data, 1) print('Image Name:{},predict:{}'.format(file, classes[pred.data.item()]))

結(jié)果如下：

4、完整代碼：

import torch.utils.data.distributed
import torchvision.transforms as transforms from PIL import Image from torch.autograd import Variable
import os from models import build_model from config import get_config
import argparse

def parse_option():
    parser = argparse.ArgumentParser('Swin Transformer Test script', add_help=False)
    parser.add_argument('--cfg', default='configs/swin_tiny_patch4_window7_224.yaml', type=str, metavar="FILE", help='path to config file', )
    parser.add_argument( "--opts", help="Modify config options by adding 'KEY VALUE' pairs. ", default=None, nargs='+',
    )

    # easy config modification
    parser.add_argument('--batch-size', type=int, help="batch size for single GPU")
    parser.add_argument('--data-path', type=str, help='path to dataset')
    parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
    parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'], help='no: no cache, ' 'full: cache all data, ' 'part: sharding the dataset into nonoverlapping pieces and only cache one piece')
    parser.add_argument('--resume', default='output/swin_tiny_patch4_window7_224/default/ckpt_epoch_1.pth', help='resume from checkpoint')
    parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps")
    parser.add_argument('--use-checkpoint', action='store_true', help="whether to use gradient checkpointing to save memory")
    parser.add_argument('--amp-opt-level', type=str, default='O1', choices=['O0', 'O1', 'O2'], help='mixed precision opt level, if O0, no amp is used')
    parser.add_argument('--output', default='output', type=str, metavar='PATH', help='root of output folder, the full path is <output>/<model_name>/<tag> (default: output)')
    parser.add_argument('--tag', help='tag of experiment')
    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
    parser.add_argument('--throughput', action='store_true', help='Test throughput only')
    parser.add_argument("--local_rank", default='0', type=int, help='local rank for DistributedDataParallel')
    args, unparsed = parser.parse_known_args()  config = get_config(args)

    return args, config  transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
classes = ("cat", "dog")




DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
_, config = parse_option()
model = build_model(config)
checkpoint = torch.load('output/swin_tiny_patch4_window7_224/default/ckpt_epoch_1.pth', map_location='cpu')
model.load_state_dict(checkpoint['model'], strict=False)
model.eval()
model.to(DEVICE)

path = 'dataset/test/' testList = os.listdir(path) for file in testList:
    img = Image.open(path + file)
    img = transform_test(img)
    img.unsqueeze_(0)
    img = Variable(img).to(DEVICE)
    out = model(img)
    # Predict
    _, pred = torch.max(out.data, 1) print('Image Name:{},predict:{}'.format(file, classes[pred.data.item()]))

總結(jié)

本文帶領(lǐng)大家學(xué)習(xí)了如何使用Swin Transformer實現(xiàn)圖像分類。通過這篇文章你學(xué)習(xí)到了Swin Transformer的環(huán)境配置和一些參數(shù)配置，學(xué)會了如何寫推理的腳本。