CSwin-T Transformer for Classification - Inference with ImageNet1K(mini version) dataset
on Projects
CSwin-T Transformer for Classification - Inference with ImageNet1K(mini version) dataset
1. dataset
2. Edit the siwn_transformer.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.helpers import load_pretrained
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from timm.models.registry import register_model
from einops.layers.torch import Rearrange
import torch.utils.checkpoint as checkpoint
import numpy as np
import time
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class LePEAttention(nn.Module):
def __init__(self, dim, resolution, idx, split_size=7, dim_out=None, num_heads=8, attn_drop=0., proj_drop=0., qk_scale=None):
super().__init__()
self.dim = dim
self.dim_out = dim_out or dim
self.resolution = resolution
self.split_size = split_size
self.num_heads = num_heads
head_dim = dim // num_heads
# NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
self.scale = qk_scale or head_dim ** -0.5
if idx == -1:
H_sp, W_sp = self.resolution, self.resolution
elif idx == 0:
H_sp, W_sp = self.resolution, self.split_size
elif idx == 1:
W_sp, H_sp = self.resolution, self.split_size
else:
print ("ERROR MODE", idx)
exit(0)
self.H_sp = H_sp
self.W_sp = W_sp
stride = 1
self.get_v = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1,groups=dim)
self.attn_drop = nn.Dropout(attn_drop)
def im2cswin(self, x):
B, N, C = x.shape
H = W = int(np.sqrt(N))
x = x.transpose(-2,-1).contiguous().view(B, C, H, W)
x = img2windows(x, self.H_sp, self.W_sp)
x = x.reshape(-1, self.H_sp* self.W_sp, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous()
return x
def get_lepe(self, x, func):
B, N, C = x.shape
H = W = int(np.sqrt(N))
x = x.transpose(-2,-1).contiguous().view(B, C, H, W)
H_sp, W_sp = self.H_sp, self.W_sp
x = x.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
x = x.permute(0, 2, 4, 1, 3, 5).contiguous().reshape(-1, C, H_sp, W_sp) ### B', C, H', W'
lepe = func(x) ### B', C, H', W'
lepe = lepe.reshape(-1, self.num_heads, C // self.num_heads, H_sp * W_sp).permute(0, 1, 3, 2).contiguous()
x = x.reshape(-1, self.num_heads, C // self.num_heads, self.H_sp* self.W_sp).permute(0, 1, 3, 2).contiguous()
return x, lepe
def forward(self, qkv):
"""
x: B L C
"""
q,k,v = qkv[0], qkv[1], qkv[2]
### Img2Window
H = W = self.resolution
B, L, C = q.shape
assert L == H * W, "flatten img_tokens has wrong size"
q = self.im2cswin(q)
k = self.im2cswin(k)
v, lepe = self.get_lepe(v, self.get_v)
q = q * self.scale
attn = (q @ k.transpose(-2, -1)) # B head N C @ B head C N --> B head N N
attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)
attn = self.attn_drop(attn)
x = (attn @ v) + lepe
x = x.transpose(1, 2).reshape(-1, self.H_sp* self.W_sp, C) # B head N N @ B head N C
### Window2Img
x = windows2img(x, self.H_sp, self.W_sp, H, W).view(B, -1, C) # B H' W' C
return x
class CSWinBlock(nn.Module):
def __init__(self, dim, reso, num_heads,
split_size=7, mlp_ratio=4., qkv_bias=False, qk_scale=None,
drop=0., attn_drop=0., drop_path=0.,
act_layer=nn.GELU, norm_layer=nn.LayerNorm,
last_stage=False):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.patches_resolution = reso
self.split_size = split_size
self.mlp_ratio = mlp_ratio
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.norm1 = norm_layer(dim)
if self.patches_resolution == split_size:
last_stage = True
if last_stage:
self.branch_num = 1
else:
self.branch_num = 2
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(drop)
if last_stage:
self.attns = nn.ModuleList([
LePEAttention(
dim, resolution=self.patches_resolution, idx = -1,
split_size=split_size, num_heads=num_heads, dim_out=dim,
qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
for i in range(self.branch_num)])
else:
self.attns = nn.ModuleList([
LePEAttention(
dim//2, resolution=self.patches_resolution, idx = i,
split_size=split_size, num_heads=num_heads//2, dim_out=dim//2,
qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
for i in range(self.branch_num)])
mlp_hidden_dim = int(dim * mlp_ratio)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, out_features=dim, act_layer=act_layer, drop=drop)
self.norm2 = norm_layer(dim)
def forward(self, x):
"""
x: B, H*W, C
"""
H = W = self.patches_resolution
B, L, C = x.shape
assert L == H * W, "flatten img_tokens has wrong size"
img = self.norm1(x)
qkv = self.qkv(img).reshape(B, -1, 3, C).permute(2, 0, 1, 3)
if self.branch_num == 2:
x1 = self.attns[0](qkv[:,:,:,:C//2])
x2 = self.attns[1](qkv[:,:,:,C//2:])
attened_x = torch.cat([x1,x2], dim=2)
else:
attened_x = self.attns[0](qkv)
attened_x = self.proj(attened_x)
x = x + self.drop_path(attened_x)
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
def img2windows(img, H_sp, W_sp):
"""
img: B C H W
"""
B, C, H, W = img.shape
img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
img_perm = img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp* W_sp, C)
return img_perm
def windows2img(img_splits_hw, H_sp, W_sp, H, W):
"""
img_splits_hw: B' H W C
"""
B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp))
img = img_splits_hw.view(B, H // H_sp, W // W_sp, H_sp, W_sp, -1)
img = img.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
return img
class Merge_Block(nn.Module):
def __init__(self, dim, dim_out, norm_layer=nn.LayerNorm):
super().__init__()
self.conv = nn.Conv2d(dim, dim_out, 3, 2, 1)
self.norm = norm_layer(dim_out)
def forward(self, x):
B, new_HW, C = x.shape
H = W = int(np.sqrt(new_HW))
x = x.transpose(-2, -1).contiguous().view(B, C, H, W)
x = self.conv(x)
B, C = x.shape[:2]
x = x.view(B, C, -1).transpose(-2, -1).contiguous()
x = self.norm(x)
return x
class SwinTransformer(nn.Module):
""" Vision Transformer with support for patch or hybrid CNN input stage
"""
def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, embed_dim=64, depths=[1,2,21,1], split_size = [1,2,7,7],
num_heads=[2,4,8,16], mlp_ratio=4., qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
drop_path_rate=0.3, hybrid_backbone=None, norm_layer=nn.LayerNorm, use_checkpoint=False):
super().__init__()
self.use_checkpoint = use_checkpoint
self.num_classes = num_classes
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
heads=num_heads
self.stage1_conv_embed = nn.Sequential(
nn.Conv2d(in_chans, embed_dim, 7, 4, 2),
Rearrange('b c h w -> b (h w) c', h = img_size//4, w = img_size//4),
nn.LayerNorm(embed_dim)
)
curr_dim = embed_dim
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depths))] # stochastic depth decay rule
self.stage1 = nn.ModuleList([
CSWinBlock(
dim=curr_dim, num_heads=heads[0], reso=img_size//4, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=split_size[0],
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[i], norm_layer=norm_layer)
for i in range(depths[0])])
self.merge1 = Merge_Block(curr_dim, curr_dim*2)
curr_dim = curr_dim*2
self.stage2 = nn.ModuleList(
[CSWinBlock(
dim=curr_dim, num_heads=heads[1], reso=img_size//8, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=split_size[1],
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[np.sum(depths[:1])+i], norm_layer=norm_layer)
for i in range(depths[1])])
self.merge2 = Merge_Block(curr_dim, curr_dim*2)
curr_dim = curr_dim*2
temp_stage3 = []
temp_stage3.extend(
[CSWinBlock(
dim=curr_dim, num_heads=heads[2], reso=img_size//16, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=split_size[2],
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[np.sum(depths[:2])+i], norm_layer=norm_layer)
for i in range(depths[2])])
self.stage3 = nn.ModuleList(temp_stage3)
self.merge3 = Merge_Block(curr_dim, curr_dim*2)
curr_dim = curr_dim*2
self.stage4 = nn.ModuleList(
[CSWinBlock(
dim=curr_dim, num_heads=heads[3], reso=img_size//32, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=split_size[-1],
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[np.sum(depths[:-1])+i], norm_layer=norm_layer, last_stage=True)
for i in range(depths[-1])])
self.norm = norm_layer(curr_dim)
# Classifier head
self.head = nn.Linear(curr_dim, num_classes) if num_classes > 0 else nn.Identity()
trunc_normal_(self.head.weight, std=0.02)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
@torch.jit.ignore
def no_weight_decay(self):
return {'pos_embed', 'cls_token'}
def get_classifier(self):
return self.head
def reset_classifier(self, num_classes, global_pool=''):
if self.num_classes != num_classes:
print ('reset head to', num_classes)
self.num_classes = num_classes
self.head = nn.Linear(self.out_dim, num_classes) if num_classes > 0 else nn.Identity()
self.head = self.head.cuda()
trunc_normal_(self.head.weight, std=.02)
if self.head.bias is not None:
nn.init.constant_(self.head.bias, 0)
def forward_features(self, x):
B = x.shape[0]
x = self.stage1_conv_embed(x)
for blk in self.stage1:
if self.use_checkpoint:
x = checkpoint.checkpoint(blk, x)
else:
x = blk(x)
for pre, blocks in zip([self.merge1, self.merge2, self.merge3],
[self.stage2, self.stage3, self.stage4]):
x = pre(x)
for blk in blocks:
if self.use_checkpoint:
x = checkpoint.checkpoint(blk, x)
else:
x = blk(x)
x = self.norm(x)
return torch.mean(x, dim=1)
def forward(self, x):
x = self.forward_features(x)
x = self.head(x)
return x
3. Edit the siwn_tiny_patch4_window7_224.yaml
MODEL:
TYPE: swin
NAME: swin_tiny_patch4_window7_224
DROP_PATH_RATE: 0.3
SWIN:
EMBED_DIM: 64
DEPTHS: [ 1, 2, 21, 1 ]
NUM_HEADS: [ 2, 4, 8, 16 ]
WINDOW_SIZE: 7
4. Edit the build.py
# --------------------------------------------------------
# Swin Transformer
# Copyright (c) 2021 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ze Liu
# --------------------------------------------------------
from .swin_transformer import SwinTransformer
from .swin_transformer_v2 import SwinTransformerV2
from .swin_transformer_moe import SwinTransformerMoE
from .swin_mlp import SwinMLP
from .simmim import build_simmim
def build_model(config, is_pretrain=False):
model_type = config.MODEL.TYPE
# accelerate layernorm
if config.FUSED_LAYERNORM:
try:
import apex as amp
layernorm = amp.normalization.FusedLayerNorm
except:
layernorm = None
print("To use FusedLayerNorm, please install apex.")
else:
import torch.nn as nn
layernorm = nn.LayerNorm
if is_pretrain:
model = build_simmim(config)
return model
if model_type == 'swin':
model = SwinTransformer(img_size=config.DATA.IMG_SIZE,
patch_size=config.MODEL.SWIN.PATCH_SIZE,
in_chans=config.MODEL.SWIN.IN_CHANS,
num_classes=config.MODEL.NUM_CLASSES,
embed_dim=config.MODEL.SWIN.EMBED_DIM,
depths=config.MODEL.SWIN.DEPTHS,
num_heads=config.MODEL.SWIN.NUM_HEADS,
#window_size=config.MODEL.SWIN.WINDOW_SIZE,
mlp_ratio=config.MODEL.SWIN.MLP_RATIO,
qkv_bias=config.MODEL.SWIN.QKV_BIAS,
qk_scale=config.MODEL.SWIN.QK_SCALE,
drop_rate=config.MODEL.DROP_RATE,
drop_path_rate=config.MODEL.DROP_PATH_RATE,
#ape=config.MODEL.SWIN.APE,
norm_layer=layernorm,
#patch_norm=config.MODEL.SWIN.PATCH_NORM,
use_checkpoint=config.TRAIN.USE_CHECKPOINT,
#fused_window_process=config.FUSED_WINDOW_PROCESS
)
elif model_type == 'swinv2':
model = SwinTransformerV2(img_size=config.DATA.IMG_SIZE,
patch_size=config.MODEL.SWINV2.PATCH_SIZE,
in_chans=config.MODEL.SWINV2.IN_CHANS,
num_classes=config.MODEL.NUM_CLASSES,
embed_dim=config.MODEL.SWINV2.EMBED_DIM,
depths=config.MODEL.SWINV2.DEPTHS,
num_heads=config.MODEL.SWINV2.NUM_HEADS,
window_size=config.MODEL.SWINV2.WINDOW_SIZE,
mlp_ratio=config.MODEL.SWINV2.MLP_RATIO,
qkv_bias=config.MODEL.SWINV2.QKV_BIAS,
drop_rate=config.MODEL.DROP_RATE,
drop_path_rate=config.MODEL.DROP_PATH_RATE,
ape=config.MODEL.SWINV2.APE,
patch_norm=config.MODEL.SWINV2.PATCH_NORM,
use_checkpoint=config.TRAIN.USE_CHECKPOINT,
pretrained_window_sizes=config.MODEL.SWINV2.PRETRAINED_WINDOW_SIZES)
elif model_type == 'swin_moe':
model = SwinTransformerMoE(img_size=config.DATA.IMG_SIZE,
patch_size=config.MODEL.SWIN_MOE.PATCH_SIZE,
in_chans=config.MODEL.SWIN_MOE.IN_CHANS,
num_classes=config.MODEL.NUM_CLASSES,
embed_dim=config.MODEL.SWIN_MOE.EMBED_DIM,
depths=config.MODEL.SWIN_MOE.DEPTHS,
num_heads=config.MODEL.SWIN_MOE.NUM_HEADS,
window_size=config.MODEL.SWIN_MOE.WINDOW_SIZE,
mlp_ratio=config.MODEL.SWIN_MOE.MLP_RATIO,
qkv_bias=config.MODEL.SWIN_MOE.QKV_BIAS,
qk_scale=config.MODEL.SWIN_MOE.QK_SCALE,
drop_rate=config.MODEL.DROP_RATE,
drop_path_rate=config.MODEL.DROP_PATH_RATE,
ape=config.MODEL.SWIN_MOE.APE,
patch_norm=config.MODEL.SWIN_MOE.PATCH_NORM,
mlp_fc2_bias=config.MODEL.SWIN_MOE.MLP_FC2_BIAS,
init_std=config.MODEL.SWIN_MOE.INIT_STD,
use_checkpoint=config.TRAIN.USE_CHECKPOINT,
pretrained_window_sizes=config.MODEL.SWIN_MOE.PRETRAINED_WINDOW_SIZES,
moe_blocks=config.MODEL.SWIN_MOE.MOE_BLOCKS,
num_local_experts=config.MODEL.SWIN_MOE.NUM_LOCAL_EXPERTS,
top_value=config.MODEL.SWIN_MOE.TOP_VALUE,
capacity_factor=config.MODEL.SWIN_MOE.CAPACITY_FACTOR,
cosine_router=config.MODEL.SWIN_MOE.COSINE_ROUTER,
normalize_gate=config.MODEL.SWIN_MOE.NORMALIZE_GATE,
use_bpr=config.MODEL.SWIN_MOE.USE_BPR,
is_gshard_loss=config.MODEL.SWIN_MOE.IS_GSHARD_LOSS,
gate_noise=config.MODEL.SWIN_MOE.GATE_NOISE,
cosine_router_dim=config.MODEL.SWIN_MOE.COSINE_ROUTER_DIM,
cosine_router_init_t=config.MODEL.SWIN_MOE.COSINE_ROUTER_INIT_T,
moe_drop=config.MODEL.SWIN_MOE.MOE_DROP,
aux_loss_weight=config.MODEL.SWIN_MOE.AUX_LOSS_WEIGHT)
elif model_type == 'swin_mlp':
model = SwinMLP(img_size=config.DATA.IMG_SIZE,
patch_size=config.MODEL.SWIN_MLP.PATCH_SIZE,
in_chans=config.MODEL.SWIN_MLP.IN_CHANS,
num_classes=config.MODEL.NUM_CLASSES,
embed_dim=config.MODEL.SWIN_MLP.EMBED_DIM,
depths=config.MODEL.SWIN_MLP.DEPTHS,
num_heads=config.MODEL.SWIN_MLP.NUM_HEADS,
window_size=config.MODEL.SWIN_MLP.WINDOW_SIZE,
mlp_ratio=config.MODEL.SWIN_MLP.MLP_RATIO,
drop_rate=config.MODEL.DROP_RATE,
drop_path_rate=config.MODEL.DROP_PATH_RATE,
ape=config.MODEL.SWIN_MLP.APE,
patch_norm=config.MODEL.SWIN_MLP.PATCH_NORM,
use_checkpoint=config.TRAIN.USE_CHECKPOINT)
else:
raise NotImplementedError(f"Unkown model: {model_type}")
return model
5. Training
python -m torch.distributed.launch --nproc_per_node 4 --master_port 12345 main.py --cfg configs/swin/swin_tiny_patch4_window7_224.yaml --data-path imagenet --batch-size 128
6. Evaluation
python -m torch.distributed.launch --nproc_per_node 1 --master_port 12345 main.py --eval --cfg configs/swin/swin_tiny_patch4_window7_224.yaml --resume checkpoints/cswin_tiny.pth --data-path imagenet
7. ver2 Inference
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.helpers import load_pretrained
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from timm.models.registry import register_model
from einops.layers.torch import Rearrange
import torch.utils.checkpoint as checkpoint
import numpy as np
import time
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class LePEAttention(nn.Module):
def __init__(self, dim, resolution, idx, split_size=7, dim_out=None, num_heads=8, attn_drop=0., proj_drop=0., qk_scale=None):
super().__init__()
self.dim = dim
self.dim_out = dim_out or dim
self.resolution = resolution
self.split_size = split_size
self.num_heads = num_heads
head_dim = dim // num_heads
# NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
self.scale = qk_scale or head_dim ** -0.5
if idx == -1:
H_sp, W_sp = self.resolution, self.resolution
elif idx == 0:
H_sp, W_sp = self.resolution, self.split_size
elif idx == 1:
W_sp, H_sp = self.resolution, self.split_size
else:
print ("ERROR MODE", idx)
exit(0)
self.H_sp = H_sp
self.W_sp = W_sp
stride = 1
self.get_v = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1,groups=dim)
self.attn_drop = nn.Dropout(attn_drop)
def im2cswin(self, x):
B, N, C = x.shape
H = W = int(np.sqrt(N))
x = x.transpose(-2,-1).contiguous().view(B, C, H, W)
x = img2windows(x, self.H_sp, self.W_sp)
x = x.reshape(-1, self.H_sp* self.W_sp, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous()
return x
def get_lepe(self, x, func):
B, N, C = x.shape
H = W = int(np.sqrt(N))
x = x.transpose(-2,-1).contiguous().view(B, C, H, W)
H_sp, W_sp = self.H_sp, self.W_sp
x = x.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
x = x.permute(0, 2, 4, 1, 3, 5).contiguous().reshape(-1, C, H_sp, W_sp) ### B', C, H', W'
lepe = func(x) ### B', C, H', W'
lepe = lepe.reshape(-1, self.num_heads, C // self.num_heads, H_sp * W_sp).permute(0, 1, 3, 2).contiguous()
x = x.reshape(-1, self.num_heads, C // self.num_heads, self.H_sp* self.W_sp).permute(0, 1, 3, 2).contiguous()
return x, lepe
def forward(self, qkv):
"""
x: B L C
"""
q,k,v = qkv[0], qkv[1], qkv[2]
### Img2Window
H = W = self.resolution
B, L, C = q.shape
assert L == H * W, "flatten img_tokens has wrong size"
q = self.im2cswin(q)
k = self.im2cswin(k)
v, lepe = self.get_lepe(v, self.get_v)
q = q * self.scale
attn = (q @ k.transpose(-2, -1)) # B head N C @ B head C N --> B head N N
attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)
attn = self.attn_drop(attn)
x = (attn @ v) + lepe
x = x.transpose(1, 2).reshape(-1, self.H_sp* self.W_sp, C) # B head N N @ B head N C
### Window2Img
x = windows2img(x, self.H_sp, self.W_sp, H, W).view(B, -1, C) # B H' W' C
return x
class CSWinBlock(nn.Module):
def __init__(self, dim, reso, num_heads,
split_size=7, mlp_ratio=4., qkv_bias=False, qk_scale=None,
drop=0., attn_drop=0., drop_path=0.,
act_layer=nn.GELU, norm_layer=nn.LayerNorm,
last_stage=False):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.patches_resolution = reso
self.split_size = split_size
self.mlp_ratio = mlp_ratio
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.norm1 = norm_layer(dim)
if self.patches_resolution == split_size:
last_stage = True
if last_stage:
self.branch_num = 1
else:
self.branch_num = 2
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(drop)
if last_stage:
self.attns = nn.ModuleList([
LePEAttention(
dim, resolution=self.patches_resolution, idx = -1,
split_size=split_size, num_heads=num_heads, dim_out=dim,
qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
for i in range(self.branch_num)])
else:
self.attns = nn.ModuleList([
LePEAttention(
dim//2, resolution=self.patches_resolution, idx = i,
split_size=split_size, num_heads=num_heads//2, dim_out=dim//2,
qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
for i in range(self.branch_num)])
mlp_hidden_dim = int(dim * mlp_ratio)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, out_features=dim, act_layer=act_layer, drop=drop)
self.norm2 = norm_layer(dim)
def forward(self, x):
"""
x: B, H*W, C
"""
H = W = self.patches_resolution
B, L, C = x.shape
assert L == H * W, "flatten img_tokens has wrong size"
img = self.norm1(x)
qkv = self.qkv(img).reshape(B, -1, 3, C).permute(2, 0, 1, 3)
if self.branch_num == 2:
x1 = self.attns[0](qkv[:,:,:,:C//2])
x2 = self.attns[1](qkv[:,:,:,C//2:])
attened_x = torch.cat([x1,x2], dim=2)
else:
attened_x = self.attns[0](qkv)
attened_x = self.proj(attened_x)
x = x + self.drop_path(attened_x)
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
def img2windows(img, H_sp, W_sp):
"""
img: B C H W
"""
B, C, H, W = img.shape
img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
img_perm = img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp* W_sp, C)
return img_perm
def windows2img(img_splits_hw, H_sp, W_sp, H, W):
"""
img_splits_hw: B' H W C
"""
B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp))
img = img_splits_hw.view(B, H // H_sp, W // W_sp, H_sp, W_sp, -1)
img = img.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
return img
class Merge_Block(nn.Module):
def __init__(self, dim, dim_out, norm_layer=nn.LayerNorm):
super().__init__()
self.conv = nn.Conv2d(dim, dim_out, 3, 2, 1)
self.norm = norm_layer(dim_out)
def forward(self, x):
B, new_HW, C = x.shape
H = W = int(np.sqrt(new_HW))
x = x.transpose(-2, -1).contiguous().view(B, C, H, W)
x = self.conv(x)
B, C = x.shape[:2]
x = x.view(B, C, -1).transpose(-2, -1).contiguous()
x = self.norm(x)
return x
class SwinTransformer(nn.Module):
""" Vision Transformer with support for patch or hybrid CNN input stage
"""
def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, embed_dim=64, depths=[1,2,21,1], split_size = [1,2,7,7],
num_heads=[2,4,8,16], mlp_ratio=4., qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
drop_path_rate=0.3, hybrid_backbone=None, norm_layer=nn.LayerNorm, use_checkpoint=False):
super().__init__()
self.use_checkpoint = use_checkpoint
self.num_classes = num_classes
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
heads=num_heads
self.stage1_conv_embed = nn.Sequential(
nn.Conv2d(in_chans, embed_dim, 7, 4, 2),
Rearrange('b c h w -> b (h w) c', h = img_size//4, w = img_size//4),
nn.LayerNorm(embed_dim)
)
curr_dim = embed_dim
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depths))] # stochastic depth decay rule
self.stage1 = nn.ModuleList([
CSWinBlock(
dim=curr_dim, num_heads=heads[0], reso=img_size//4, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=1,
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[0], norm_layer=norm_layer)
])
self.merge1 = Merge_Block(curr_dim, curr_dim*2)
curr_dim = curr_dim*2
self.stage2 = nn.ModuleList(
[CSWinBlock(
dim=curr_dim, num_heads=heads[1], reso=img_size//8, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=2,
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[np.sum(depths[:1])+0], norm_layer=norm_layer),
CSWinBlock(
dim=curr_dim, num_heads=heads[1], reso=img_size//8, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=2,
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[np.sum(depths[:1])+1], norm_layer=norm_layer)
])
self.merge2 = Merge_Block(curr_dim, curr_dim*2)
curr_dim = curr_dim*2
temp_stage3 = []
temp_stage3.extend(
[CSWinBlock(
dim=curr_dim, num_heads=heads[2], reso=img_size//16, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=split_size[2],
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[np.sum(depths[:2])+i], norm_layer=norm_layer)
for i in range(depths[2])])
self.stage3 = nn.ModuleList(temp_stage3)
self.merge3 = Merge_Block(curr_dim, curr_dim*2)
curr_dim = curr_dim*2
self.stage4 = nn.ModuleList(
[CSWinBlock(
dim=curr_dim, num_heads=heads[3], reso=img_size//32, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, qk_scale=qk_scale, split_size=7,
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[np.sum(depths[:-1])+0], norm_layer=norm_layer, last_stage=True)
])
self.norm = norm_layer(curr_dim)
# Classifier head
self.head = nn.Linear(curr_dim, num_classes) if num_classes > 0 else nn.Identity()
trunc_normal_(self.head.weight, std=0.02)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
@torch.jit.ignore
def no_weight_decay(self):
return {'pos_embed', 'cls_token'}
def get_classifier(self):
return self.head
def reset_classifier(self, num_classes, global_pool=''):
if self.num_classes != num_classes:
print ('reset head to', num_classes)
self.num_classes = num_classes
self.head = nn.Linear(self.out_dim, num_classes) if num_classes > 0 else nn.Identity()
self.head = self.head.cuda()
trunc_normal_(self.head.weight, std=.02)
if self.head.bias is not None:
nn.init.constant_(self.head.bias, 0)
def forward_features(self, x):
B = x.shape[0]
x = self.stage1_conv_embed(x)
for blk in self.stage1:
if self.use_checkpoint:
x = checkpoint.checkpoint(blk, x)
else:
x = blk(x)
for pre, blocks in zip([self.merge1, self.merge2, self.merge3],
[self.stage2, self.stage3, self.stage4]):
x = pre(x)
for blk in blocks:
if self.use_checkpoint:
x = checkpoint.checkpoint(blk, x)
else:
x = blk(x)
x = self.norm(x)
return torch.mean(x, dim=1)
def forward(self, x):
x = self.forward_features(x)
x = self.head(x)
return x
