本文轉載自:AI人工智能初學者

Rethinking BiSeNet For Real-time Semantic Segmentation
論文:https://arxiv.org/abs/2104.13188
代碼(即將開源):
https://github.com/MichaelFan01/STDC-Seg
超強的實時語義分割新網絡:STDC,表現SOTA!性能優於BiSeNetV2、SFNet等網絡,其中一版本性能高達97.0FPS/76.8%mIoU,代碼即將開源!作者單位:美團
1 簡介
BiSeNet已被證明在實時分割two-stream網絡中是有效的。但是,其添加額外path以對空間信息進行編碼的原理很耗時,並且由於缺少任務專用設計,因此從預訓練任務(例如圖像分類)中借用的主干可能無法有效地進行圖像分割。
為了解決這些問題,作者提出了一種新穎且有效的結構,即通過消除結構冗余來緩解以上的耗時問題(Short-Term Dense Concatenate network)。具體來說,本文將特征圖的維數逐漸降低,並將特征圖聚合起來進行圖像表征,形成了STDC網絡的基本模塊。在decoder中,提出了一個Detail Aggregation module將空間信息的學習以single-stream方式集成到low-level layers中。最后,將low-level features和deep features融合以預測最終的分割結果。

圖1 性能對比圖
在Cityscapes上,在1080Ti上以250.4FPS的速度在測試集上實現了71.9%的mIoU,比最新方法快45.2%,在97.0FPS的情況下達到76.8%的mIoU,同時可以推斷出更高分辨率的圖像。
2 問題動機
為了更快的推理以達到實時的推理,很多的Researcher提出了很多的方法,也使用了很多輕量化的backbone,諸如DFANet、BiSeNetV1等模型;但是這些輕量化的Backbone主要是針對圖像分類的設計,可能針對特定的語義分割任務性能並不是那么的好;
除此之外輸入圖像的分辨率似乎對於推理速度也有比較大的影響,但這樣也很容易導致模型忽略目標邊界和小物體的細節。

圖2 BiSeNet與本文所提方法
如圖2(a)所示,為了解決以上問題,BiSeNet采用multi-path框架將low-level details和high-level semantics結合起來。然而,添加一條額外的path來獲取low-level特征是很費時的,同時Auxiliary path也往往缺乏low-level信息的引導。
2.1 編碼階段
針對上訴問題作者設計了一種新的網絡,目的是更快的推理速度、更具有可解釋的結構和的現有方法對比更具競爭的性能。
首先,設計了一個新的結構Short-Term Dense Concatenate module模塊(STDC模塊),以通過少量的參數來獲得Variant Scalable Receptive Fields。
其次,將STDC模塊集成到U-net體系結構中,形成STDC Network,這一操作極大地提高了語義分割任務網絡的性能。

圖3 General STDC網絡架構
如圖3所示,將多個連續層的Feature maps連接起來,每個層對輸入的圖像/特征在不同的尺度和各自的域進行編碼,從而實現多尺度特征表示。為了加快速度,逐步減小層的卷積核大小,同時在分割性能上的損失可以忽略不計。
2.2 解碼階段
在解碼階段,如圖2(b)所示,采用Detail Guidance來引導low-level layers進行空間細節的學習,而不是使用額外的Path。
首先,利用Detail Aggregation模塊生成詳細的Ground-truth。
然后,利用binary cross-entropy loss和dice loss優化細節信息,將細節信息作為一種約束信息在推理時不需要。
最后,融合low-level layers信息和semantic信息預測語義分割結果。
3 Proposed Method
3.1 Design of Encoding Network
3.1.1 Short-Term Dense Concatenate Module

圖4 STDC模塊布局
圖4(b)和(c)展示了STDC模塊的布局。具體來說,每個模塊被分成幾個子模塊,用表示第塊的運算。因此,第塊的輸出如下:

其中和分別為第塊的輸入和輸出。包括卷積層、BN層和ReLU,是卷積層的內核大小。
在STDC模塊中,第一個塊的卷積核大小為1,其余的設置為3。給定STDC模塊輸出通道數為N;除了最后一個卷積層的卷積核數與前一個卷積層的卷積核數相同,其余第塊中卷積層的卷積核數均為。
在圖像分類任務中,More Channels in Higher Layers是一種常見的做法。但在語義分割任務中側重於Scalable Receptive Field和Multi-Scale。Low-level layers需要足夠的通道來編碼Receptive Field較小的更細粒度的信息,而Receptive Field較大的high-level layers更注重high-level信息的引導,與Low-level layers設置相同的通道可能會導致信息冗余。
Downsample只在Block2中使用。為了獲得更加豐富特征信息,作者選擇通過skip-path cat和作為STDC模塊的輸出。在cat之前,要先通過3×3平均池化操作,將STDC模塊中不同塊的response maps降采樣到相同的空間大小,STDC模塊的最終輸出為:

式中,為STDC模塊輸出,F為方法中的融合操作,是所有塊的特征映射。考慮到效率,作者采用cat作為fusion操作。在本文中作者使用4個block構建STDC模塊。

STDC Module的優勢:
- 通過幾何級數的方式精心調整塊的卷積核的大小,從而顯著降低了計算復雜度。
- STDC模塊的最終輸出由所有塊cat而成,保留了scalable respective fields和multi-scale information。
3.1.2 復雜度分析
給定輸入通道維度,輸出通道維度,則STDC模塊parameter number為:

如式3所示,STDC模塊的parameter number由預定義的輸入輸出通道維度決定,而block個數對parameter number的影響較小。特別是當n達到最大值時,STDC模塊的參數數幾乎保持不變,僅由M和n定義。
3.1.3 Network Architecture

圖5 Network Architecture

除輸入層和預測層外,還包括6個階段:
通常,Stage-1~5對輸入的空間分辨率分別進行了2個下采樣,Stage-6通過一個ConvX、一個全局平均池化層和兩個全連接層輸出預測logits。
Stage-1和Stage-2通常被是用於提取外觀特征的low-level layers。為了追求效率,在Stage-1和Stage-2中每個階段只使用一個卷積塊,根據經驗證明這是有效的。Stage-3、4、5中的STDC模塊的數量在網絡中是經過仔細調整的。在這些Stage中,每個Stage的第1個STDC模塊使用stride=2的卷積塊進行下采樣,緊接着Stage的保持空間分辨率不變。
3.2 Design of Decoder
3.2.1 Segmentation Architecture
首先,使用預訓練STDC網絡作為編碼器的Backbone,並采用BiSeNet的context path對context information進行編碼。
如圖3(a)所示,使用Stage-3,4,5分別生成了降樣比為1/8,1/16,1/32的feature map。然后使用全局平均池化生成具有large receptive field的全局上下文信息。
采用U-shape結構對全局特征進行上采樣,並將其與編碼階段的后2個階段(Stage-4和Stage-5)的特征進行組合。
在BiSeNet之后使用Attention Refine模塊來細化每2個階段的組合特性。
在最終的語義分割預測中采用了BiSeNet中的Feature Fusion模塊,將從Stage-3得到的1/8降采樣特征與解碼器得到的1/8降采樣特征進行融合。作者認為這2個Stage的特征處於不同的特征表征層次。來自編碼Stage的特征保留了豐富的細節信息,來自解碼器的特征由於全局池化層的輸入而包含了上下文信息。
具體來說,Seg Head包含一個3×3 Convon-Bn-Relu算子,然后進行1×1卷積,得到輸出維數N(類的個數)。采用Cross-Entry Loss-OHEM優化語義分割學習任務。
3.2.2 Detail Guidance of Low-level Features

圖5 STDC Segmentation network
與相同下采樣比的backbone low-level layers(Stage-3)相比,spatial path可以編碼更多的spatial細節,如邊界、角等。在此基礎上,本文提出了一個細節引導模塊來引導low-level layers以單流方式學習空間信息。作者將細節預測建模為二值分割任務。
-
首先,利用Laplacian算子對分割后的Ground-truth進行細節映射detail map Ground-truth(如圖5(c)所示)。在Stage-3中插入detail Head生成detail feature map。
-
然后,使用detail ground-truth作為detail feature map的引導來指導low-level layers學習spatial detail feature。如圖6(d)所示,經過detail引導的feature map比圖6(c)所示的結果可以編碼更多的空間細節。

圖6 細節特征對比圖
- 最后,將學習到的細節特征與解碼器深度塊的上下文特征融合,進行分割預測。
1、Detail Ground-truth Generation
通過detail Aggregation模塊從語義分割ground truth中生成binary detail ground-truth,如圖5(c)中虛線藍框所示。這種運算可以通過2-D拉普拉斯核卷積和可訓練的卷積來實現。
-
首先,使用如圖5(e)所示的Laplacian算子生成不同步幅的soft thin detail feature map以獲取多尺度細節信息。
-
然后,我們將細節特征映射上采樣到原始大小,並融合一個可訓練的11 - 1卷積來進行動態重加權。
-
最后,利用邊界和角點信息采用閾值0.1將預測細節轉化為最終的binary detail ground-truth。
2、Detail Loss
由於detail pixels的數量遠少於非detail pixels,很明顯detail預測是一個類不平衡問題。由於加權交叉熵總是導致粗糙的結果,於是作者采用binary cross-entropy和dice loss聯合優化細節學習。dice loss衡量的是預測map和ground-truth之間的重疊。此外,它對前景/背景像素的數量不敏感,這意味着它可以緩解類不平衡的問題。因此,對於高度H、寬度W的預測detail圖,detail loss 公式如下:

其中為預測的細節,為相應的細節ground-truth。為binary cross-entropy loss,為dice loss:

如圖5(b)所示使用Detail Head生成Detail map,該map引導淺層對空間信息進行編碼。Detail Head包括一個3×3 con-bn-relu算子,然后是一個1×1卷積來得到輸出的Detail map。
在實驗中Detail Head被證明是有效的增強特征表示。請注意,此分支將在推斷階段被丟棄。因此,這種側信息可以在不耗費任何推理代價的情況下輕松地提高分割任務的准確性。
4. 實驗
4.1 ImageNet實驗

本文與最新的輕量級backbone進行對比,之類與常用的ResNet50進行一下比較可以看出來STDC2精度上比ResNet50高出1.1個百分點,但是速度是ResNet50的2.5倍,FLOPs是ResNet50的約1/3.
4.2 語義分割實驗

從上表可以看出來,盡管mIoU在Cityscapes數據集不是最高的,甚至還不如比較火的DANet,但是速度卻是最快的,居然達到了250FPS。只能說實在是香!!!
4.3 可視化對比

論文PDF和代碼下載
后台回復:STDC,即可下載上述論文PDF
后台回復:CVPR2021,即可下載CVPR 2021論文和代碼開源的論文合集
后台回復:Transformer綜述,即可下載最新的兩篇Transformer綜述PDF
重磅!圖像分割交流群成立
掃碼添加CVer助手,可申請加入CVer-圖像分割 微信交流群,方向已涵蓋:目標檢測、圖像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超分辨率、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感圖像、行為識別、視頻理解、圖像融合、圖像檢索、論文投稿&交流、PyTorch和TensorFlow等群。
一定要備注:研究方向+地點+學校/公司+昵稱(如圖像分割+上海+上交+卡卡),根據格式備注,可更快被通過且邀請進群

▲長按加小助手微信,進交流群
▲點擊上方卡片,關注CVer公眾號
