Swin Transformer:用Transformer實現CNN多尺度操作

文本是關于Swin Transformer基礎知識的了解

論文：https://arxiv.org/pdf/2103.14030

項目：https://github. com/microsoft/Swin-Transformer.

實現一個Swin Transformer：Swin Transformer模型具體代碼實現-CSDN博客

Swin Transformer mlp版本實現：Swin Transformer 變體1：使用MLP代替多頭注意力機制-CSDN博客

Swin Transformer v2版本實現：Swin Transformer變體2：Swin Transformerv2模型-CSDN博客

原理總覽：

1、總述：

Swin Transformer和Vision Transformer很像，他們都是將NLP領域的Transformer架構引入計算機視覺領域期待大力出奇跡的產物，不同的是，Vision Transformer由于它固定的尺度（16x16大小的patch），它并不適合密集預測，比較適合于圖像分類任務，對于目標檢測、語義分割這種任務就有點困難，而且最主要的問題是，ViT使用的是全局注意力機制，這種注意力的最大問題是計算成本和內存消耗太大。而Swin transformer的目標不是做一類任務，它想做視覺任務的骨干模型。

Swin Transformer 做出的第一大改進就是在窗口內計算注意力，這樣的話就把復雜度從與圖片大小的平方關系改進到了線性關系，因為窗口內計算注意力的計算量相比全局注意力的計算量差距很大，很多個窗口完全可以并行在batch中計算。

第二大改進就是引入了移動窗口，因為Swin Transformer 使用窗口注意力，那么就有一個問題，不同窗口之間的信息怎么交互呢，作者引入了窗口移動機制，使用兩個block層，第一個block層中使用現有的窗口做自注意力計算，之后移動窗口，具體來說就是將patch矩陣沿著右下角移動window_size//2個單位，之后在第二個block中計算新的窗口中的注意力計算，在計算完成之后進行反向移動操作，將之前的移位操作復原。

第三大改進就是學習CNN的pooling機制，實現了層次化的多尺度建模，這樣的話有助于學習不同尺度的特征，對密集預測十分有幫助，這點在論文的實驗結果部分表現得很清楚。

Swin Transformer的這些特性使其與廣泛的視覺任務兼容，包括圖像分類（在ImageNet-1K上達到87.3%的top-1準確率）和密集預測任務，如目標檢測（在COCO測試集上達到58.7的box AP和51.1的mask AP）和語義分割（在ADE20K驗證集上達到53.5的mIoU）。其性能超越了之前的最先進水平，COCO上提高了+2.7 box AP和+2.6 mask AP，ADE20K上提高了+3.2 mIoU，展示了基于Transformer的模型作為視覺骨干網絡的潛力。層次化設計和移動窗口方法也證明了對全MLP架構的益處。

2、總體架構：

Swin Transformer總體來說分為4個階段：

首先，對于輸入的圖像數據，需要先進行Patch Partition操作，也就是分塊，這里的patch大小為4x4，相比于ViT的patch大小16x16，Swin Transformer的每個patch可以保留更細節的信息。如上圖，(H//4)*(W//4)表示的是patch數量，48表示的是3x4x4，也就是一個patch的三個通道的所有像素的拼接，以此作為每個patch的特征，之后經過嵌入，將48映射為C（tiny和small模型中默認為96，big和large模型默認為192）,這里所有的操作其實在代碼實現中都是通過一個卷積層實現的，具體可以參考ViT的實現。

2.1、stage 1：

上圖架構是Swin Transformer-T模型的架構圖，他的第一階段有2個block層，在Swin Transformer中，2個block作為一組同時出現，所以每個階段的block數量都為偶數：

在第一個階段中，對于輸入的數據【B，H//4*W//4，C】：

在第一個block中不進行窗口移動，只進行窗口自注意力計算，我們可以詳細分析一下這中間的特征形狀變化：對于輸入數據形狀：【B，H//4*W//4，C】，首先將一維序列展開為二位序列：【B，H//4，W//4，C】，之后劃分窗口：【B*num_windows，windows_size，windows_size，C】，其中windows_size是窗口大小，默認為7，num_windows為窗口數量=（H//4//windows_siz，W//4//windows_size），之后根據每個block的注意力頭的數量重塑形狀為：【B*num_windows，head，windows_size*windows_size，C//head】，其中，windows_size*windows_size就是一個窗口內的序列長度，C//head就是每個元素的特征，這樣的話，就把原本ViT的全局注意力計算變為了窗口注意力計算（ViT：【B，head，（H//16）*（W//16）+1，C//head】），這里需要注意，在計算完計算力得分之后，還有一個相對位置偏移需要相加，這一點將在本章的第5節講解。計算完注意力，需要將窗口合并，即形狀變為【B，H//4，W//4，C】，之后將數據形狀重塑回到原始形狀【B，H//4*W//4，C】即可進入下一個block。

在第二個block中，需要進行窗口的移動，這個操作在劃分窗口前進行，只需要調用torch.roll函數即可。在注意力計算階段，由于窗口發生了變化，相對應的注意力計算區域也發生了變化，為了保證窗口數量不變以及不計算不應該計算位置的注意力得分，這里對于 $QK^{T}$ 的結果要進行掩碼處理，具體操作在本章的第6節講解。進行完注意力計算之后，需要進行窗口合并，即形狀變為【B，H//4，W//4，C】，然后將之前移動的窗口還原，之后將數據形狀重塑回到原始形狀【B，H//4*W//4，C】。

2.2、stage 2：

stage 2第一個操作是降采樣操作，也就是patch_merging，這個將在本章的第三節進行講解，降采樣的結果就是將圖像尺寸減半，將通道加倍，這個操作是在模仿CNN中的池化操作，經過該操作輸入數據變為【B，H//8*W//8，2C】。

stage 2 的Swin Transformer block的流程和stage1 完全一致，就是輸入數據形狀不一樣。

2.3、stage 3：

stage 3第一個操作和stage 2一樣，都是降采樣把圖像尺寸減半，將通道加倍，這個操作是在模仿CNN中的池化操作，經過該操作輸入數據變為【B，H//16*W//16，4C】。

stage 3 的Swin Transformer block的流程和stage1 完全一致，就是輸入數據形狀不一樣，同時stage 3有6個block。

2.4、stage 4：

stage 4第一個操作和stage 2一樣，都是降采樣把圖像尺寸減半，將通道加倍，這個操作是在模仿CNN中的池化操作，經過該操作輸入數據變為【B，H//32*W//32，8C】。

stage 4 的Swin Transformer block的流程和stage1 完全一致，就是輸入數據形狀不一樣。

2.5、輸出處理：

stage 4的輸出經過平均池化和壓縮之后，得到輸出結果，形狀為：[B, num_features]，之后經過Swin Transformer head的映射，輸出結果[B, num_classes]。

3、patch_merging：

該模塊是模型的降采樣模塊，是把圖像尺寸減半，將通道加倍，這個操作是在模仿CNN中的池化操作。下面的圖可以簡要描述這種操作：

上圖中的1，2，3，4代表的是應該在分割后歸屬的圖的標號而不是真實值，按照上圖所示的邏輯，將原始數據分割為4部分，每一部分形狀為?【B，H//2，W//2，C】，在C通道維度拼接為【B，H//2，W//2，4C】，之后將4C維度映射為2C：【B，H//2，W//2，2C】。

4、窗口自注意力機制和全局自注意力機制運算量：

$\Omega (MSA)=4hwC^{2}+2(hw)^{2}C$

$\Omega (WMSA)=4hwC^{2}+2M^{2}hwC$ ?

其中，h和w是圖像的高和寬，C是通道數量，M是窗口大小。

我們知道全局多頭自注意力機制的運算包括4個矩陣映射、2個矩陣乘法，其中，矩陣映射是K、Q、V（h*w*C）與C*C的矩陣的相乘，運算量為 $3hwC^{2}$ ，當計算完注意力分數之后，要將attn經過映射輸出，這部分也是h*w*C與C*C的矩陣的相乘，運算量為 $hwC^{2}$ 。兩個矩陣乘法是 $QK^{T}$ 和 $AV$ ，他們都是h*w*C與C*h*w矩陣乘法，運算量都為 $(hw)^{2}C$ ，所以全局注意力運算量總和為 $4hwC^{2}+2(hw)^{2}C$ 。

窗口多頭自注意力機制的運算也包括4個矩陣映射、2個矩陣乘法，其中，矩陣映射是K、Q、V（M*M*C）與C*C的矩陣的相乘，運算量為 $3M^{2}C^{2}$ ，當計算完注意力分數之后，要將attn經過映射輸出，這部分也是M*M*C與C*C的矩陣的相乘，運算量為 $M^{2}C^{2}$ 。兩個矩陣乘法是 $QK^{T}$ 和 $AV$ ，他們都是M*M*C與C*M*M矩陣乘法，運算量都為 $M^{4}C$ ，所以一個窗口注意力運算量總和為 $4M^{2}C^{2}+2M^{4}C$ ,總的窗口個數是（h//M）*（w//M），相乘得到窗口注意力運算量總和為 $4{hw}C^{2}+2M^{2}hwC$ 。

5、相對位置偏移：

為什么要采用相對位置？ViT的結果表明使用相對位置和絕對位置對最后的分類結果影響不大，但是Swin Transformer它的用途不局限于圖像分類，在目標檢測、語義分割等密集預測任務中，使用相對位置比使用絕對位置的效果要好很多，這一點作者做了消融實驗，結果如下：

Swin Transformer的相對位置偏移表大小為【（2M-1）*（2M-1），head_num】，其中，2M-1是窗口的高和寬的相對位置范圍，即如果M=3，那么（-2，-1，0，1，2）就是相對位置范圍。相對位置偏移表是一個可學習的參數表。

然后使用窗口的高和寬堆疊形成二維網格，之后壓縮形成【2，M*M】的網格坐標，之后計算第一個M*M坐標和第二個M*M坐標每一對之間的相對坐標差，得到【2，M*M，M*M】，之后重塑形狀為【M*M，M*M，2】并給每一個元素加上M-1使得所有元素均不為負，這是因為這些坐標差之后要作為index在相對位置偏移表中取值，之后將最后一維求和得到【M*M，M*M】，將其作為相對位置索引表。

當計算完 $QK^{T}$ 之后，將相對位置索引表中所有元素對應于相對位置偏移表中的head_num維向量取出并重塑形狀為【head_num，M*M，M*M】。

我們可以回憶一下 $QK^{T}$ 的形狀：【B，head_num，patch_num，patch_num】，其中patch_num=M*M，那么我們給【head_num，M*M，M*M】擴充第零維之后利用廣播就可以實現相對位置偏移和 $QK^{T}$ 的相加。

6、窗口移動：

我們知道Swin Transformer的特殊點就在于其可以移動的窗口，但是移動窗口就會出現一個問題：

一個原始特征圖如果有4個窗口，沿著右下角移動window_size//2個單位后會發現出現了9個窗口，我們知道如果窗口高和寬不一致，我們無法使用注意力機制，這時最簡單的辦法就是填充，將高和寬不一致的窗口用0填充至高和寬一致，這樣做很簡單，但是相應的計算量就上升了，這對于我們來說是不利的，所以我們換一種思路，引入掩碼操作：

上圖中，每種顏色覆蓋范圍代表一個邏輯上形式上的窗口（4個），每種數字標識一個實際的窗口（9個），在同一個顏色的邏輯窗口內可能會存在不同的實際窗口，我們可以給一個邏輯窗口內的所有實際窗口做差值，使得掩碼中同一個實際窗口中的元素為0，非該實際窗口中的元素為-100，計算注意力時使用的是邏輯上的窗口，這樣能保證窗口數量不變。把這個掩碼值加入到 $QK^{T}$ 結果中，做softmax操作就可以不需要計算注意力的地方置為0。