打破ViT SOTA壟斷！SegNeXt：卷積注意力機制重奪語義分割的勝利高地（NeurIPS 22）

NeurIPS2022： http://github.com/Visual-Attention-Network/SegNeXt/blob/main/resources/paper.pdf

code： http://github.com/Visual-Attention-Network/SegNeXt

arxiv： http://arxiv.org/abs/2209.08575

Vision Transformer的"降維打擊"導致多個CV領域SOTA方案均被ViT方案主導，語義分割同樣不例外。本文對已有成功分割方案進行了重審視並發現了幾個有助於效能提升的關鍵成分，進而促使我們設計了一種新型的卷積注意力架構方案SegNeXt。在多個主流語義分割資料集上，SegNeXt大幅改善了其效能。比如，在Pascal VOC2012測試集上，SegNeXt憑藉僅需EfficientNet-L2+NAS-FPN的十分之一引數量取得了90.6%mIoU指標。平均來講，在ADE20K資料集上，SegNeXt比其他SOTA方案平均提高2.0%mIoU，且計算量相當或更少。

本文出發點

本文對語義分割領域代表性方案(DeepLabV3+, HRNet, SETR, SegFormer)進行重審視，總結出成功的語義分割方案應具有的幾點關鍵屬性(可參考上表)：

• 採用強骨幹網路作為編碼器；

• 多尺度資訊互動；

• 空域注意力；

• 低計算複雜度。

基於上述考量，本文對卷積注意力設計進行了重思考並提出了一種簡單而有效的編碼器-解碼器架構SegNeXt。不同於已有Transformer方案，SegNeXt對Transformer-Convolution Encoder-Decoder架構進行了逆轉，即對編碼器模組採用傳統卷積模組設計但引入了多尺度卷積注意力，對解碼器模組採用了Hamberger(自注意力的一種替代方案)進一步提取全域性上下文資訊。因此，SegNeXt能夠從區域性到全域性提取多尺度上下文資訊，能在空域與通道維度達成自適應性，能從底層到高層進行資訊聚合。上圖給出了Cityscape與ADE20K資料集上所提方案與標杆方案的計算量與效能對比。

本文方案

接下來，我們將對SegNeXt所用到的定製化骨幹網路(其實就是擴充套件版VAN)與Decoder部分進行介紹。

Convolutional Encoder

延續已有方案，本文在Encoder部分同樣採用了金字塔架構，每個構成模組採用了類似ViT的結構，但不同之處在於：本文並未使用自注意力，而是設計一種多尺度卷積注意力模組MSCA(見下圖，其實就是一種多尺度版VAN)。

如上圖所示，MSCA由三部分構成：

• depth-wise 卷積：用於聚合區域性資訊

• 多分支depth-wise卷積：用於捕獲多尺度上下文資訊

• 卷積：用於在通道維度進行相關性建模

其中，卷積的輸出將作為注意力權值對MSCA的輸入進行重加權。MSCA可以描述為如下公式：

在實現方面，多分支depth-wise方面，作者採用兩個depth-wise strip卷積近似標準depth-wise卷積。這樣做有兩個好處：

• strip卷積更輕量，可以進一步降低計算量；

• 對於類strip目標(如行人、電線杆)，strip卷積可以作為標準卷積的互補提取類strip卷積。

上表給出了通過堆疊MSCA而得到的不同MSCAN骨幹資訊以及SegNeXt架構資訊。此外，主要注意的是：MSCAN的每個模組採用的是BN，而非LN。關於MSCA的實現code參考如下。

class AttentionModule(BaseModule):
    def __init__(self, dim):
        super().__init__()
        self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
        self.conv0_1 = nn.Conv2d(dim, dim, (1, 7), padding=(0, 3), groups=dim)
        self.conv0_2 = nn.Conv2d(dim, dim, (7, 1), padding=(3, 0), groups=dim)

        self.conv1_1 = nn.Conv2d(dim, dim, (1, 11), padding=(0, 5), groups=dim)
        self.conv1_2 = nn.Conv2d(dim, dim, (11, 1), padding=(5, 0), groups=dim)

        self.conv2_1 = nn.Conv2d(dim, dim, (1, 21), padding=(0, 10), groups=dim)
        self.conv2_2 = nn.Conv2d(dim, dim, (21, 1), padding=(10, 0), groups=dim)
        self.conv3 = nn.Conv2d(dim, dim, 1)

    def forward(self, x):
        u = x.clone()
        attn = self.conv0(x)

        attn_0 = self.conv0_1(attn)
        attn_0 = self.conv0_2(attn_0)

        attn_1 = self.conv1_1(attn)
        attn_1 = self.conv1_2(attn_1)

        attn_2 = self.conv2_1(attn)
        attn_2 = self.conv2_2(attn_2)
        attn = attn + attn_0 + attn_1 + attn_2

        attn = self.conv3(attn)

        return attn * u

Decoder

常規語義分割模型的骨幹往往在ImageNet上預訓練得到，為捕獲高階語義資訊，通常需要一個Decoder模組。本文則對上述三種簡單Decoder架構進行了探索：

• Figure3-a，源自SegFormer的解碼器，它是一種純MLP架構；

• Figure3-b，常被CNN方案使用，如ASPP、PSP、DANet等；

• Figure3-c，本文采用的解碼器，它採用輕量型Hamberger模組對後三個階段的特性進行聚合以進行全域性上下文建模。

需要注意的是，SegFormer的解碼器對Stage1到Stage4的特徵進行聚合，而本文方案則僅對Stage2-Stage4的特徵進行聚合。這是因為：

• SegNeXt的Encoder採用了卷積架構，使得Stage1部分特徵包含過多底層資訊，進而導致其會影響語義分割效能。

• 對Stage1部分特徵進行處理會帶來過多的計算負載。

上圖為Hamberger的架構示意圖，它採用矩陣分解方式進行全域性空域資訊建模。關於Hamberger的介紹，作者在文中並未進行描述，該部分內容源自作者被ICLR2021接收的文章《Is Attention Better Than Matrix Decomposition》，關於該部分的介紹查閱作者的詳細解讀：

• Enjoy Hamburger：注意力機制比矩陣分解更好嗎？(I)，知乎連結：http://zhuanlan.zhihu.com/p/369769485

• Enjoy Hamburger：注意力機制比矩陣分解更好嗎？(II)，知乎連結：http://zhuanlan.zhihu.com/p/369855045

• Enjoy Hamburger：注意力機制比矩陣分解更好嗎？(III)，知乎連結：http://zhuanlan.zhihu.com/p/370410446

話說作者的理論功底是真的強！強烈建議各位同學跟隨作者的思路去領略一下Hamberger背後的故事。從github上的介紹來看該文在3月份就已初步完成，不過當時是採用的是VAN+Hamberger組合，而本文則是MSCAN+Hamberger的組合。當然，MSCAN是VAN的多尺度版。

本文實驗

上述圖表給出了不同資料集上不同語義分割方案的效能對比，從中可以看到：

• 在ADE20K資料集上，相比Mask2Former+Swin-T，SegNeXt-L以3.3mIoU指標勝出，且引數量與計算量相當；相比SegFormer-B2，SegNeXt-B以2.0mIoU指標勝出，且僅需56%計算量。

• 在Cityscape資料集上，相比SegFormer-B2，SegNeXt-B以1.6mIoU指標勝出，且僅需40%引數量。

更多實驗結果可參考如上三個表，總而言之，SegNeXt在語義分割任務上取得了新的SOTA結果。值得一提的是，在Cityscape測試集+實時分割方面，所提方案達成了新的SOTA指標78.0mIoU@25fps(3090 RTX GPU)。

# 極市平臺簽約作者 #

投稿方式：

新增小編微信Fengcall（微訊號：fengcall19），備註： 姓名-投稿

△長按新增極市平臺小編