NeurIPS 2022 | 百度提出超快Transformer分割模型RTFormer,180FPS+81mIOU


最近,基于
Transformer的網(wǎng)絡(luò)在語義分割方面取得了令人印象深刻的成果。然而,對于實時語義分割,由于Transformer的計算機制非常耗時,純CNN方法仍然是該領(lǐng)域的主流。本文提出了
RTFormer,一種用于實時語義分割的高效雙分辨率Transformer,與基于CNN的模型相比,它在性能和效率之間取得了更好的平衡。為了在類似GPU的設(shè)備上實現(xiàn)高推理效率,RTFormer利用具有線性復(fù)雜性的GPU Friendly Attention,并放棄了多頭機制。此外,作者發(fā)現(xiàn),通過傳播從低分辨率分支學(xué)到的高層次知識,跨分辨率注意力更有效地為高分辨率分支收集全部上下文信息。
對主流基準(zhǔn)測試的大量實驗證明了提出的
RTFormer的有效性,它在Cityscapes、CamVid和COCOStuff上達(dá)到了最先進(jìn)的水平,并在ADE20K上顯示出良好的結(jié)果。
1、簡介
語義分割是一項基本的計算機視覺任務(wù),通常作為自主駕駛、移動應(yīng)用、機器人感知等領(lǐng)域的關(guān)鍵感知模塊。隨著這些應(yīng)用的發(fā)展,實時語義分割的需求越來越強烈。現(xiàn)有的實時分割方法主要側(cè)重于利用CNN架構(gòu),包括通過手工設(shè)計高效主干網(wǎng)絡(luò)和解碼器,以及探索神經(jīng)架構(gòu)搜索方法,以在準(zhǔn)確性和效率之間找到更好的平衡。到目前為止,這些優(yōu)秀的作品已經(jīng)取得了顯著的進(jìn)步。
近年來,視覺Transformer因其強大的視覺識別能力而備受關(guān)注。在繼承它們的基礎(chǔ)上,提出了一系列基于Transformer的體系結(jié)構(gòu),并在一般的語義分割任務(wù)中表現(xiàn)出很好的性能。與基于CNN的網(wǎng)絡(luò)相比,這些基于Transformer的架構(gòu)的主要區(qū)別在于大量使用自注意力,而自注意力擅長捕獲遠(yuǎn)程上下文信息,這在語義分割中是必不可少的。因此,作者認(rèn)為注意力結(jié)構(gòu)在實時語義分割任務(wù)中也應(yīng)該有效。
但到目前為止,只有少數(shù)作品探索了注意力在這一領(lǐng)域的應(yīng)用,目前的SOTA水平仍然由基于CNN的架構(gòu)主導(dǎo)。作者認(rèn)為,在實時環(huán)境中應(yīng)用注意力的主要障礙可能來自以下兩個方面。
一是現(xiàn)有的大多數(shù)注意力類型的計算特性對于類似 GPU的設(shè)備來說都不是很容易推理的,例如二次復(fù)雜度和多頭機制。二次復(fù)雜度在處理高分辨率特征時引入了巨大的計算負(fù)擔(dān),特別是在語義分割等密集預(yù)測任務(wù)中。盡管有幾部作品縮小了key和value的大小,但二次復(fù)雜性的特性仍然存在。而多頭機制將矩陣乘法拆分為多個組,這使得注意力操作在類似GPU的設(shè)備上耗費時間,類似于執(zhí)行組卷積的情況。另一個原因是,僅關(guān)注高分辨率特征圖本身可能不是用高級別語義信息捕獲遠(yuǎn)程上下文的最有效方法,因為來自高分辨率特征圖的單個特征向量的感受野有限。
本文提出了一種新的Transformer塊,名為RTFormer塊,如圖2所示,其目的是在具有Transformer的類似GPU的設(shè)備上實現(xiàn)性能和效率之間的更好平衡。對于低分辨率分支,采用了新提出的GPU Friendly Attention,該注意力源自External Attention。它繼承了External Attention的線性復(fù)雜性特性,并通過丟棄矩陣乘法操作中的通道分裂,緩解了類GPU設(shè)備的多頭機制的弱點。相反,它擴大了外部參數(shù)的數(shù)量,并將External Attention提出的雙范數(shù)運算中的第二個歸一化分解為多個組。這使得GPU Friendly Attention能夠在一定程度上保持多頭機制的優(yōu)勢。對于高分辨率分支,采用交叉分辨率注意力,而不是僅在高分辨率特征本身內(nèi)進(jìn)行注意力。此外,與多分辨率融合的并行公式不同,作者將兩個分辨率分支排列成階梯布局。因此,通過輔助從低分辨率分支中學(xué)習(xí)到的高級別全局上下文信息,可以更有效地增強高分辨率分支。

基于提出的RTFormer塊,構(gòu)建了一個新的實時語義分割網(wǎng)絡(luò),稱為RTFormer。為了了解足夠的局部上下文,RTFormer仍然在早期階段使用卷積塊,并在最后兩個階段放置RTFormer塊。通過大量的實驗發(fā)現(xiàn)RTFormer可以更有效地利用全局上下文,并比以前的工作取得更好的折衷效果。圖1顯示了RTFormer與CamVid上其他方法的比較。最后,我們將RTFormer的貢獻(xiàn)總結(jié)為以下三個方面:
提出了一種新的 RTFormer塊,它在類GPU設(shè)備上實現(xiàn)了語義分割任務(wù)的性能和效率之間的更好平衡。提出了一種新的網(wǎng)絡(luò)體系結(jié)構(gòu) RTFormer,它可以充分利用全局上下文,通過深度利用注意力而不損失效率來改進(jìn)語義分割。RTFormer在Cityscapes、CamVid和COCOStuff上達(dá)到了最先進(jìn)的水平,在ADE20K上表現(xiàn)良好。此外,它為實時語義分割任務(wù)的實踐提供了一個新的視角。
2、本文方法
2.1、RTFormer block
RTFormer塊是繼承多分辨率融合范式的雙分辨率模塊。與之前的工作相比,RTFormer塊由兩種類型的注意力及其前饋網(wǎng)絡(luò)組成,并按階梯式布局排列,如圖2所示。在低分辨率分支中使用GPU Friendly Attention來捕獲高級全局上下文。在高分辨率分支中,引入了交叉分辨率注意力,將從低分辨率分支中學(xué)習(xí)到的高級別全局上下文廣播到每個高分辨率像素,并使用階梯式布局將低分辨率分支的更具代表性的特征引入交叉分辨率注意力。

1、GPU-Friendly Attention
通過比較現(xiàn)有的不同類型的注意力發(fā)現(xiàn)External Attention(EA)由于其令人滿意的線性復(fù)雜度特性,可以成為在類似GPU的設(shè)備上執(zhí)行的潛在選擇,GPU-Friendly Attention(GFA)就是從中派生出來的。因此,在詳細(xì)說明GFA之前,首先說明EA。讓表示一個輸入特征,其中是元素(或圖像中的像素)的數(shù)量,是特征維數(shù),那么EA的公式可以表示為:

其中是可學(xué)習(xí)的參數(shù),是參數(shù)維,DN是雙重規(guī)范化操作。EA的多頭版本可以解釋為:

其中,是頭數(shù),而是的第個頭。如圖3的上半部分所示,多頭機制生成個注意力圖,以提高EA的容量,這使得矩陣乘法被分成幾個組,類似于組卷積。雖然EA對不同的頭使用共享的和,這可以大大加快計算速度,但分裂矩陣乘法仍然存在。

為了避免類似GPU的設(shè)備由于多頭機制而延遲減少,作者提出了一種簡單有效的GPU-Friendly Attention。它是由方程式1所表示的基本External Attention演變而來的,可公式化為:

其中,GDN表示分組雙重歸一化,它將原始雙重歸一化的第二次歸一化分解為組,如圖3左下部分所示。從方程3中可以發(fā)現(xiàn)GFA有兩個主要改進(jìn)。
一方面,它可以集成矩陣乘法,這對于類似GPU的設(shè)備非常友好。得益于此,可以將外部參數(shù)的大小從擴大到。因此,可以調(diào)整更多參數(shù)以提高性能。
另一方面,它利用分組雙重歸一化在一定程度上保持了多頭機制的優(yōu)勢。為了直觀理解,可以認(rèn)為
GFA還生成了個不同的注意力圖,用于捕獲Token之間的不同關(guān)系,但計算相似性需要更多的特征元素,所有注意力圖都有助于最終輸出。
2、Cross-resolution Attention
多分辨率融合已被證明對密集預(yù)測任務(wù)是有效的。而對于多分辨率架構(gòu)的設(shè)計,可以直觀地將GFA獨立應(yīng)用于不同分辨率的分支,并在執(zhí)行卷積模塊或注意力模塊后交換特征。但在高分辨率分支中,像素更關(guān)注局部信息,而非高級別的全局上下文。因此,作者認(rèn)為直接關(guān)注高分辨率特征圖來學(xué)習(xí)全局背景是不夠有效的。為了更有效地獲取全局上下文,作者提出了一種交叉分辨率注意力,旨在充分利用從低分辨率分支中學(xué)習(xí)到的高層語義信息。如圖2所示,與GFA不同,高分辨率分支采用交叉分辨率注意力來收集全局背景。交叉分辨率注意力的計算表示為:

其中分別表示高分辨率分支和低分辨率分支上的特征映射,是一組矩陣運算,包括 splitting、permutation以及reshaping,表示高分辨率分支的特征維數(shù)。值得說明的是,下文中表示為交叉特征的特征圖是通過由池化層和卷積層組成的函數(shù)從計算得出的。的空間大小表示從生成的Token數(shù)。
在實驗上,只在注意力圖的最后一個軸上使用softmax進(jìn)行歸一化,因為當(dāng)key和value不是外部參數(shù)時,單個softmax比雙重歸一化性能更好。特別是,為了在類似GPU的設(shè)備上進(jìn)行快速推理,這里也放棄了多頭機制。
3、Feed Forward Network
在以前的基于Transformer的分割方法中,前饋網(wǎng)絡(luò)(FFN)通常由兩個MLP層和一個深度3×3卷積層組成,其中深度3×3卷積用于補充位置編碼或增強局部性。此外,兩個MLP層將隱藏維度擴展為輸入維度的2到4倍。這種類型的FFN可以用相對較少的參數(shù)獲得更好的性能。但在應(yīng)該考慮類GPU設(shè)備上的延遲的情況下,FFN的典型結(jié)構(gòu)不是很有效。為了平衡性能和效率,在RTFormer塊的FFN中采用了兩個3×3卷積層,而沒有進(jìn)行維數(shù)擴展。與典型的FFN配置相比,它顯示了更好的結(jié)果。
2.2、RTFormer

1、Backbone Architecture
為了提取高分辨率特征圖所需的足夠的局部信息,將卷積層與RTFormer塊相結(jié)合來構(gòu)造RTFormer。具體來說,讓RTFormer從一個由兩個3×3卷積層組成的干塊開始,用幾個連續(xù)的基本殘差塊組成前兩個階段。然后,從第3階段開始,使用雙分辨率模塊,實現(xiàn)高分辨率和低分辨率分支之間的特征交換。對于最后三個階段的高分辨率分支,特征stride保持為8不變,而對于低分辨率分支,則特征stride分別為16、32、32。特別是,將雙分辨率模塊安排為階梯式布局,以借助低分辨率分支的輸出增強高分辨率特征的語義表示。最重要的是,用提出的RTFormer塊構(gòu)造了stage4和stage5,如圖2所示,用于有效的全局上下文建模,而stage3仍然由基本殘差塊組成。
2、Segmentation Head
對于RTFormer的分割頭,在低分辨率輸出特征之后添加了一個DAPPM模塊。然后將DAPPM的輸出與高分辨率特征融合,得到stride為8的輸出特征圖。最后,將該輸出特征傳遞到像素級分類頭,用于預(yù)測密集語義標(biāo)簽。分類頭由3×3卷積層和1×1卷積層組成,隱藏特征維數(shù)與輸入特征維數(shù)相同。
3、Instantiation

用RTFormer Slim和RTFormer Base實例化了RTFormer的架構(gòu),詳細(xì)配置記錄在表1中。對于通道數(shù)和塊數(shù),每個數(shù)組包含5個元素,分別對應(yīng)于5個階段。特別是,具有兩個數(shù)字的元素對應(yīng)于雙分辨率級。例如,64/128表示通道數(shù)對于高分辨率分支為64,對于低分辨率分支為128。而1/2表示高分辨率分支的基本卷積塊數(shù)為1,低分辨率分支為2。
值得注意的是,塊編號數(shù)組中的最后兩個元素表示RTFormer塊的數(shù)量,對于RTFormer Slim和RTFormer Base,它們都是1。RTFormer Slim和RTFormer Base的交叉特征空間尺寸分別設(shè)置為64(8×8)和144(12×12)。
3、實驗
3.1、消融實驗
1、不同形式注意力的對比

為了驗證提出的注意力的有效性,將RTFormer塊中使用的注意力替換為不同類型和組合。如表5a所示,給出了多頭自注意力、多頭external attention、GPU-Friendly attention 和交叉分辨注意力的不同組合的結(jié)果。
例如,“GFA+CA”意味著在低分辨率分支中使用GFA,在高分辨率分支中則使用CA。此外,通過M=d×r調(diào)整了多頭外部注意中的超參數(shù)M,其中r是一個縮減率??梢园l(fā)現(xiàn),GPU-Friendly attention 優(yōu)于所有多頭外部注意設(shè)置,并且在M=d時比最佳注意設(shè)置更快,同時,GPU-Friendly attention 比多頭自注意力更有效,性能相當(dāng)。
這表明,在類似GPU的設(shè)備上,GPU-Friendly attention比多頭自注意力和多頭external attention在性能和效率之間取得更好的平衡。當(dāng)引入交叉分辨率注意力時,性能會進(jìn)一步提高,而FPS僅下降不到2。
2、不同形式的FFN的對比

表5b顯示了由兩個MLP層和一個3×3深度卷積層組成的典型FFN的結(jié)果,以及提出的包含兩個3×3卷積層的FFN。結(jié)果表明,提出的FFN不僅在mIoU上,而且在FPS上都優(yōu)于典型的FFN。這表明提出的FFN更適合于應(yīng)該考慮類GPU設(shè)備延遲的情況。
3、分組雙重標(biāo)準(zhǔn)化中組數(shù)的影響

研究了在使用GPU Friendly Attention對兩個分支進(jìn)行分組雙重標(biāo)準(zhǔn)化的情況下,分組數(shù)目的影響。表5c顯示了不同配置的結(jié)果。例如,“8+2”表示在低分辨率分支中使用8個組,在高分辨率分支中則使用2個組。特別是,當(dāng)組的數(shù)目設(shè)置為1時,分組的雙重歸一化會退化為原始的雙重歸一化。在這里,當(dāng)組數(shù)為8和2時,可以獲得最佳mIoU,這說明分組的雙重歸一化比原始的雙重歸一化性能更好。值得注意的是,改變分組雙重歸一化中的組數(shù)并不影響推理效率,這使得GPU Friendly Attention能夠在組數(shù)較大時保持較高的FPS。
4、交叉特征的空間大小對交叉分辨率注意的影響

作者還研究了交叉分辨率注意力中交叉特征的空間大小,包括應(yīng)用6×6、8×8和12×12。如表5d所示,根據(jù)FPS和mIoU之間的權(quán)衡,RTFormer Slim的交叉特征的8×8空間大小是最佳的。在一定程度上,這表明與高分辨率特征尺寸相近的交叉特征的空間尺寸是合適的,因為RTFormer Slim的高分辨率特征維數(shù)為64,等于8×8。
3.2、泛化性分析

表3顯示了關(guān)于ADE20K的結(jié)果。RTFormer Base實現(xiàn)了42.1%的卓越mIoU和71.4FPS,優(yōu)于所有其他方法。例如,與DDRNet-23-Slim相比,RTFormer Slim實現(xiàn)了更好的mIoU 36.7%,并且保持了幾乎相同的速度。圖6顯示了ADE20K驗證集的定性結(jié)果。與DDRNet-23相比,RTFormer顯示了更好的細(xì)節(jié)和上下文信息??傊@些結(jié)果表明,RTFormer在廣義場景中的實時語義分割方面也顯示出非常有前景的性能。在COCOStuff上,如表4所示,RTFormer Base以143.3 FPS的速度達(dá)到35.3 mIoU,以相當(dāng)?shù)耐评硭俣瘸^DDRNet-23約3%,并創(chuàng)下了新的一流水平。
3.3、SOTA對比



3.4、分類實驗對比

4、參考
[1].RTFormer: Efficient Design for Real-Time Semantic Segmentation with Transformer.
5、推薦閱讀
即插即用 | CNN與Transformer都通用的Trick,即插即漲點即提速!
目標(biāo)檢測改進(jìn) | 如何使用IOU改進(jìn)自注意力以提升Sparse RCNN目標(biāo)檢測性能
重參巔峰 | 你喜歡的RepVGG其實也是有缺陷的,RepOpt才是重參的巔峰
