基于Transformer 的道路場景分割算法研究

魏鵬磊，雷菊陽

（上海工程技術大學 機械與汽車工程學院，上海 201620）

0 引言

常見的基于深度卷積神經(jīng)網(wǎng)絡的模型，如FCN［1］、Deeplab v3［2］、SegNet［3］等在傳統(tǒng)語義分割任務上有著很好的效果，但是對于城市道路場景的分割仍然難以達到理想的狀態(tài)。隨著Transformer在NLP 領域表現(xiàn)優(yōu)異性能之后，越來越多的人嘗試將其應用在CV 領域，并取得了可觀進步。繼ViT［4］之后，出現(xiàn)了很多Transformer 運用在CV 各個任務上的工作，而Swin Transformer［5］是第一個備受青睞的可以在下游任務中使用的純Transformer 結構的方式，但卻有著如下缺點：參數(shù)量過大、顯存占用高、訓練時間長。究其原因，下游任務，如語義分割是高密度預測任務，對于分割精度要求很高，從而使得訓練參數(shù)量巨大，增加了訓練成本。故本文在SwinT的基礎上改進了網(wǎng)絡結構，可以明顯加快訓練速度，也可以很好地定位分割邊界。其次，針對特征信息學習不充分問題，傳統(tǒng)做法為通過設置不同參數(shù)的卷積層或池化層，先提取到不同尺度的特征圖，再將這些特征圖送入網(wǎng)絡做融合。但是由于圖像金字塔的多尺度輸入，在計算時需要保存大量梯度，故對硬件的要求很高。而本次研究是將網(wǎng)絡進行多尺度訓練，在測試階段進行多尺度融合，這樣可減少參數(shù)和內(nèi)存占用，且由于引入多尺度信息，可以更好地定位分割邊界，提高了網(wǎng)絡性能。

1 相關工作

1.1 數(shù)據(jù)集

本文使用Pascal VOC 2012 擴增數(shù)據(jù)集做基礎研究，Cityscapes 數(shù)據(jù)集做進一步驗證。

Pascal VOC 挑戰(zhàn)賽是一個世界級的計算機視覺挑戰(zhàn)賽。Pascal VOC 挑戰(zhàn)賽整體上可分為如下幾類：圖像分類、目標檢測、目標分割、行為識別等。在Pascal VOC 數(shù)據(jù)集中主要包含20 個目標類別和1 個背景類別。

對于圖像語義分割，Pascal VOC 2012 中共有訓練集圖像1 464 張、驗證集圖像1 449 張、測試集圖像1 456張，但是對于語義分割，特別是基于Transformer 骨干網(wǎng)絡而言，擁有大量的數(shù)據(jù)是很有必要的，所以本文使用了Pascal VOC 的擴增數(shù)據(jù)集，共有訓練集圖像10 582張。

另外，在語義分割中對應的標注圖像（.png）用PIL的Image.open（）函數(shù)讀取時，默認是P 模式（調(diào)色板模式），即一個單通道的圖像。在背景處的像素值為0，目標邊緣處用的像素值為255，目標區(qū)域根據(jù)目標類別的類別索引信息進行填充，如圖1所示，人對應的目標索引是15，所以目標區(qū)域的像素值用15 填充。具體調(diào)色板信息見表1。

圖1 P 模式下的標簽圖Fig.1 Label map in P mode

表1 不同類別的索引值Tab.1 Index values for different categories

Cityscapes 數(shù)據(jù)集于2016 年發(fā)布，在自動駕駛領域是權威且熱門的語義分割數(shù)據(jù)集之一，該數(shù)據(jù)集含有國外50 個道路場景的高分辨率圖像，其中精細標注圖像共有5 000張，粗略標記圖像共19 998張，為保證能夠最大限度地獲取充足的數(shù)據(jù)信息。本文使用含粗略標注和精細標注數(shù)據(jù)集24 998張，共分為建筑、行人、天空等19 個類別。

1.2 數(shù)據(jù)預處理

基于Transformer 網(wǎng)絡架構相對于深度卷積神經(jīng)網(wǎng)絡更容易出現(xiàn)過擬合現(xiàn)象，除需對網(wǎng)絡中的模型結構進行優(yōu)化外，擁有大量的數(shù)據(jù)也能夠減少過擬合的發(fā)生，故考慮對圖像進行預處理。本文的數(shù)據(jù)擴增操作是在OpenCV 上完成的，包括對圖像進行-10°～10°的旋轉、隨機裁剪crop_size的0.5～2倍、隨機水平翻轉以及模糊圖像等操作。

2 方法

本文算法是由2 個路徑組成的，分別是編碼器提取路徑與解碼器提取路徑。其中，編碼器塊是在Swin Transformer 的基礎上改進后得到的，不僅加快了訓練速度，而且也緩解了過擬合。解碼器塊中的Prediction Head 則是基于ASPP+模塊，考慮通過利用跳躍連接以及捷徑分支優(yōu)化模塊結構，使其可以更好地解決目標多尺度的問題。具體來說，是以改進后的Swin Transformer 模型SwinLab 為骨干網(wǎng)絡，再對ASPP 模塊進行優(yōu)化，并構建模塊ASPP+，使ASPP+可以多尺度理解上下文信息的能力。整體模型通過3 個階段構建不同大小的特征圖，且又在SwinT 的基礎上剔除掉Patch Partition 和Linear Embeding 模塊，并添加1 個和后2 個階段同樣的Patch Merging 層進行下采樣。網(wǎng)絡總體模型結構如圖2 所示。

圖2 網(wǎng)絡總體模型結構Fig.2 Overall model structure of the network

編碼塊是有2 個結構，一個使用了W-MSA 結構，另一個使用SW-MSA 結構。一般情況下，這2 個結構是成對使用的，先使用W-MSA 結構，而后使用SW-MSA 結構，具體編碼器模型如圖3、圖4 所示。

圖3 編碼器結構圖Fig.3 Encoder structure diagram

圖4 MLP 結構圖Fig.4 MLP structure diagram

解碼塊包括ASPP+模塊和Prediction Head 模塊。ASPP+在ASPP 的基礎上摒棄了膨脹系數(shù)為36的空洞卷積層，并采用自適應平均池化層，即共有4個并行分支，分別為1 個1×1 卷積層、3 個3×3 空洞卷積層，以及1 個自適應全局平均池化層，該層目的是可以增加1 個全局上下文信息。其中，使用concat方法對4 個并行分支進行拼接之前，先利用自注意力機制對不同分支獲得的信息進行注意力處理，這樣有利于不同特征信息的融合，而虛線部分的捷徑分支則使用1×1 卷積核進行維度處理。對于Prediction Head 模塊來說，得到ASPP+模塊的輸出后，添加一個跳躍連接殘差模塊［6］，其后續(xù)接一個Layer Norm層，再通過一個1×1 卷積層來融合信息。Prediction Head 通過雙線性插值的方法還原輸入圖像的尺寸大?。?］，網(wǎng)絡模型細節(jié)如圖5 所示。

圖5 解碼器結構圖Fig.5 Decoder structure diagram

3 實驗結果分析

在Pascal VOC2012 數(shù)據(jù)集和Cityscapes 數(shù)據(jù)集上的分割效果如圖6、圖7 所示。圖6、圖7中，從（a）到（d）分別是原圖、標簽、DeepLabv3 預測圖以及SwinLab 預測圖。

圖6 Pascal VOC2012 數(shù)據(jù)集Fig.6 Pascal VOC2012 dataset

圖7 Cityscapes 數(shù)據(jù)集Fig.7 Cityscapes dataset

4 結束語

針對道路場景識別任務，本文提出了一種基于Transformer 的SwinLab 模型架構。該網(wǎng)絡架構增強了網(wǎng)絡在多尺度下多類別分割時的魯棒性，同時使用不同的采樣比例與感受野提取特征，使其可以在多個尺度上捕獲上下文信息。實驗結果表明，基于Transformer 構建的SwinLab 模型網(wǎng)絡相比于傳統(tǒng)基于深度卷積神經(jīng)網(wǎng)絡的語義分割模型，效果，及性能均獲提升，雖不及SOTA，但在Pascal VOC2012 數(shù)據(jù)集上mIoU可達80.1，在Cityscapes 數(shù)據(jù)集上也有不錯的效果。除此之外，本文重點關注的訓練速度也得到了顯著改善，對于后續(xù)的研究有著實際參考意義。另外，本文使用的顯卡為單張英偉達最新3090顯卡，再加上網(wǎng)絡性能指標與實驗環(huán)境有較大的相關性，故理論上本文所構建網(wǎng)絡的性能仍然有較大可提升的空間。