CSWin Transfomer：超越Swin Transformer的網(wǎng)絡來了

近期，微軟亞研院繼Swin Transformer之后又推出了CSWin Transformer。和Swin Transformer一樣，CSWin Transformer也是一種local self-attention網(wǎng)絡，相比Swin的方形window self-attention，CSWin采用的是十字形（cross-shaped）window self-attention，這使得CSWin Transformer的建模能力更強，在分類和檢測等任務上也超過Swin Transformer，其中CSWin-L在語義分割數(shù)據(jù)集ADE20K上達到了SOTA：55.7 mIoU，超過Swin-L的53.5（不過目前微軟提出的無監(jiān)督訓練模型BEiT-L已經(jīng)再次刷新了榜單：57.0 mIoU）。CSWin Transformer和Swin Transformer一樣采用金字塔結(jié)構(gòu)，共包括4個stage，各個stage的特征圖大小分別是原圖的1/4，1/8，1/16和1/32。CSWin Transformer主要有三個重要的改進：Overlapping Patch Embedding，Cross-Shaped Window Self-Attention和Locally-Enhanced Positional Encoding。

Overlapping Patch Embedding

PVT和Swin Transformer等較早的金字塔模型中patch embedding是沒有overlap的，patch size為的patch embedding操作上等價于stride和kernel size均為的卷積，所以模型開始的patch embedding就是一個stride為4的4x4卷積，而后面各個stage間的patch merging就是一個stride為2的2x2卷積。但是隨后的CvT和PVTv2都采用overlapping patch embedding，這個變動是對性能有提升的。因此，CSWin Transformer也采用overlapping patch embedding：開始的patch embedding采用stride為4的7x7卷積，而后面各個stage間的patch merging采用stride為2的3x3卷積：

# patch embedding
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)
)

# patch merging
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

注意，這里的卷積都需要包含zero padding來保持和原來一樣的輸出大小。

Cross-Shaped Window Self-Attention

CSWin Transformer最核心的部分就是cross-shaped window self-attention，如下所示，首先將self-attention的mutil-heads均分成兩組，一組做horizontal stripes self-attention，另外一組做vertical stripes self-attention。所謂horizontal stripes self-attention就是沿著H維度將tokens分成水平條狀windows，對于輸入為HxW的tokens，記每個水平條狀window的寬度為，那么共產(chǎn)生個windows，每個window共包含個tokens；而vertical stripes self-attention就是沿著W維度將tokens分成豎直條狀windows，同樣地會產(chǎn)生個windows，每個window的tokens量為。具體的劃分窗口代碼和Swin transformer一樣，通過設定window的寬度和長度來實現(xiàn)兩組attention：

# 對于水平attention，H_sp=sw, W_sp=W
# 對于豎直attention，H_sp=H, W_sp=sw
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

兩組self-attention是并行的，完成后將tokens的特征concat在一起，這樣就構(gòu)成了CSW self-attention，最終效果就是在十字形窗口內(nèi)做attention，CSW self-attention的感受野要比常規(guī)的window attention的感受野更大。用公式表示的話就是：

可以得到CSWin attention的計算復雜度為，普通的window attention的計算復雜度是和成正比的，而global attention的計算復雜度和的平方成正比的，而CSWin attention的計算復雜度介于兩者之間。另外一點是CSWin transformer中不同的stage采用不同的，前面的stage采用較小的，而后面的stage采用較大，這其實也是漸進式地擴大感受野。默認4個stage的分別設為1, 2, 7, 7。CSWin attention的代碼實現(xiàn)如下所示：

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 # sw
        self.mlp_ratio = mlp_ratio
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.norm1 = norm_layer(dim)
  
        # 最后一個階段，實際上執(zhí)行的是global attention
        if self.patches_resolution == split_size:
            last_stage = True
        if last_stage:
            self.branch_num = 1 # 只有一個分支
        else:
            self.branch_num = 2 # 兩個分支，分別執(zhí)行兩組attention
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(drop)
        # 最后一個階段，就只有一個window，不需要再分成兩組
        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)]) # idx區(qū)分兩組attention
        

        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]) # 一半heads執(zhí)行水平attention
            x2 = self.attns[1](qkv[:,:,:,C//2:]) # 另外一半heads執(zhí)行豎直attention
            attened_x = torch.cat([x1,x2], dim=2) # concat在一起
        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

從代碼實現(xiàn)可以看到兩點，首先是對最后一個stage，由于輸入為已經(jīng)為7x7（輸入圖像為224x224），而也是7，那么其實只有一個window，就等于在做global attention，也就沒必要再分成兩組了。而對于前面3個stage，其實CSWin attention是分成兩支的，分別做兩種attention，雖然兩者是相對獨立的，但是也是分開做的，主要有兩個原因，一是兩種attention的窗口數(shù)量不一定相同（當H和W不相等時），二是兩種attention的positional encoding也是不同的。另外CSWin attention和早期的Sequential Axial很類似，不過后者。論文中也對各種attention機制做了對比實驗，無論是分類，檢測還是分割，CSWin attention都是更勝一籌（這里CSWin采用non-overlapping patch embedding以及Swin的positional encoding來減少其它因素的干擾）：

Locally-Enhanced Positional Encoding

CSWin Transformer采用的也是一種relative positional encoding（RPE)，不過不同于常規(guī)RPE將位置信息加在attention的計算上，這里考慮將位置信息直接施加在上，如下所示：

考慮到的計算量較大，這里用一個depth-wise convolution（3x3卷積）來替換，這其實就主要考慮局部位置信息了，論文稱這種位置編碼為locally-enhanced positional encoding (LePE)：

由于是卷積，所以LePE可以接受任意輸入大小，對下游任務如檢測和分割比較友好，其具體實現(xiàn)如下：

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
        # 最后一個stage
        if idx == -1:
            H_sp, W_sp = self.resolution, self.resolution
        elif idx == 0: # 水平attention
            H_sp, W_sp = self.resolution, self.split_size
        elif idx == 1: # 豎直attention
            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
       
        # LePE
        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

論文中也對各種位置編碼方式做了對比，可以看到LePE在各個任務上效果均最好：

CSWin Transformer

CSWin Transformer的網(wǎng)絡設置如下，也包括4個不同大小的模型，其主要區(qū)別在channels和各個stages的depth：在ImageNet分類上，CSWin Transformer要優(yōu)于Swin Transformer和Twins等模型：在COCO實例分割任務上，CSWin Transformer的AP也要優(yōu)于Swin Transformer和Twins等模型：在語義分割ADE20K數(shù)據(jù)集上，最終的CSWin-L的mIoU達到了55.7：

小結(jié)

相比Swin Transformer，CSWin Transformer更進了一步，這也是local attention網(wǎng)絡在CV任務上的勝利。其實同期微軟團隊還有另外一篇論文Focal Self-attention for Local-Global Interactions in Vision Transformers也取得了較好的性能，但是效果稍微比CSWin Transformer差一些（Focal-L在ADE20K數(shù)據(jù)集上達到了55.4）。

參考

1. CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows
2. PVTv2: Improved Baselines with Pyramid Vision Transformer
3. Focal Self-attention for Local-Global Interactions in Vision Transformers