Abstract
本文研究了利用大型網絡訓練小型語義分割網絡的知識蒸餾策略。我們從簡單的像素式精餾方案開始, 該方案應用于圖像分類的精餾方案, 并分別對每個像素進行知識精餾。我們進一步提出將結構化知識從大型網絡提煉成小型網絡, 其動機是語義分割是一個結構化預測問題。我們研究兩種結構化蒸餾方案: (i) 對精餾, 蒸餾成對的相似性, (ii) 整體蒸餾, 使用 GAN 提取整體知識。通過在雪鐵龍、坎維德和 ADE20K 三個場景解析數據集上的大量實驗, 證明了我們的知識蒸餾方法的有效性
1. Introduction
語義分割是預測輸入圖像中每個像素的類別標簽的問題。它是計算機視覺中的一項基本任務, 具有許多實際應用, 如自動駕駛、視頻監控、虛擬現實等。自完全卷積神經網絡 (FCNs) 發明以來, 深部神經網絡一直是語義分割的主要解決方案 [38]。隨后的方法, 例如 DeepLab [5, 6, 7, 48]、PSPNet [56]、OCNet [50]、RefineNet[23] 和 DenseASPP [46], 在分割精度方面取得了顯著提高, 通常采用了繁瑣的模型和昂貴的計算。
近年來, 由于移動設備應用的需要, 具有模型尺寸小、計算成本輕、分割精度高的神經網絡越來越受到人們的關注。目前的大部分工作都是專門設計輕量級網絡, 用于分割或借用分類網絡的設計, 例如 ENet [31]、ESPNet [31]、ERFNet [34] 和 ICNet [55]。本文的重點在于緊湊型分割網絡, 重點是利用繁瑣的網絡對緊湊型網絡進行訓練, 以提高分割精度。
我們研究了知識蒸餾策略, 該策略已被驗證在分類任務 [15, 35] 中的有效性, 用于訓練緊湊型語義分割網絡。作為一個簡單的方案, 我們簡單地將分割問題看作是許多單獨的像素分類問題, 然后直接將知識精餾方案應用到像素級。這個簡單的方案, 我們稱之為pixel-wise distillation, 將相應像素的類概率從繁瑣的網絡 (教師) 轉移到緊湊的網絡 (學生)。
考慮到語義分割是一個結構化預測問題, 我們提出了結構信息的結構解餾, 并采用了一對式精餾和整體蒸餾兩種方案來傳遞結構信息。基于配對的蒸餾方案是由廣泛研究的對馬爾可夫隨機場框架 [22] 來加強空間標記連續性, 目的是加強空間標記的連續性, 其目的是使從緊湊型網絡和繁瑣的網絡計算的像素之間的對等相似性對齊。
整體蒸餾方案的目的是在緊湊型分割網絡生成的分割映射和繁瑣的分割之間, 對不具有像素化和對向式精餾特征的高階一致性進行對齊網絡。我們采用了對抗性訓練方案, 鼓勵從緊湊分割網絡生成的分割映射的整體嵌入, 而不是與繁瑣分割網絡的輸出區分開來。
為此, 我們優化了一個目標函數, 該函數將傳統的多級交叉熵損失與蒸餾項結合起來。本文的主要貢獻可概述如下。
·我們研究知識蒸餾策略, 以訓練準確的緊湊型語義分割網絡
·我們提出了兩種結構化的知識蒸餾方案, pair-wise和整體蒸餾, enforcing pair-wise and high-order consistency between the outputs of the compact and cumbersome segmentation networks 。
·我們通過在三個基準數據集上改進最近開發的最先進的緊湊型分割網絡 espnet、Mobiletv2-plus 和 ResNet18 來展示我們的方法的有效性: Cityscapes [10]、CamVid [4] 和 ADE20K [58], 這是如圖1所示
2. Related Work
語義分割。深層卷積神經網絡自創先爭優以來一直是語義分割的主要解決方案, fully-convolutionalnetwork [38], DeConvNet [30], U-Net [36]。為了提高網絡能力和分割性能, 開發了各種方案 [47]。例如, 更強大的主干網絡 (如 Google 網 [39]、Resnet [14] 和 Densenet [17]) 表現出更好的分割性能。通過膨脹卷積 [5、6、7、48] 或多徑細化網絡 [23] 來改進分辨率, 從而顯著提高性能。利用多尺度的上下文, 例如, 膨脹的卷積 [48], 在 PSPNet [56] 中的金字塔池模塊, 在 DeepLab [6] 中的atrous [6]spatial pyramid pooling, 對象上下文 [50], 也有利于分割。Lin 等人 [24] 將深層模型與結構化輸出學習相結合, 用于語義分割。
除了繁瑣的網絡用于高度精確的分割之外, 由于需要真正的應用程序 (如移動應用程序), 高效的分割網絡也越來越吸引人們的興趣。大多數工作都集中在輕量級網絡設計上, 通過分解技術加速卷積操作。ENet [31] 的靈感來自于 [40], 集成了幾個加速因子, 包括多分支模塊、早期特征圖分辨率下采樣、小解碼器尺寸、濾波器張量分解等。SQ [41] 采用SqueezeNet [18] fire模塊和并行膨脹卷積層, 以實現高效分割。ESPNet [28] 提出了一個有效的空間金字塔, 它是基于過濾分解技術: point-wise卷積和空間金字塔的膨脹卷, 以取代標準卷積。高效的分類網絡, 如 MobileNet [16]、ShuffleNet [54] 和 IGCNet [53], 也用于加速分割。此外, ICNet (圖像級聯網絡) [55] 利用處理低分辨率圖像的效率和高分辨率圖像的高推理質量, 實現效率和準確性之間的權衡
知識蒸餾。知識蒸餾 [15] 是一種將知識從繁瑣的模型轉變為緊湊型模型的方法, 以提高緊湊型網絡的性能。將繁瑣模型產生的類概率作為訓練緊湊型模型 [2, 15, 42] 或傳輸中間特征圖 [35, 51] 的軟目標, 將其應用于圖像分類。還有其他應用, 包括物體檢測 [21], 行人重新識別 [9] 等。最近開發的語義分割應用 [45] 與我們的方法有關。它主要分別提取每個像素的類概率 (如我們的像素化蒸餾) 和每個局部補丁的標簽的中心差異 (在 [45] 中稱為局部關系)。相反, 我們專注于提煉結構化知識: pairwise 蒸餾, 它轉移所有像素對之間的關系, 而不是local patch中的關系 [45], 以及整體蒸餾, 它傳遞捕獲的整體知識高階信息。
對抗性學習。生成對抗性網絡 (gans) 在文本生成 [43, 49] 和圖像合成 [12, 20] 中得到了廣泛的研究。條件版本 [29] 成功地應用于圖像到圖像的翻譯, 包括樣式傳輸 [19]、圖像繪制 [32]、圖像著色 [26] 和文本到圖像 [33]。對抗性學習的想法也被采用在姿勢估計 [8] 鼓勵人的姿態估計結果不被區分從地面真相;和語義分割 [27], 鼓勵估計分割映射不與地面真值圖區分開來。[27] 中的一個挑戰是發電機的連續輸出和離散真實標簽之間的不匹配, 使 GAN 中的鑒別器的成功非常有限。不同于 [27], 在我們的方法中, 使用的 GAN 沒有這個問題, 因為鑒別器的基本真理是教師網絡的日志, 這是真正的價值。我們使用對抗性學習來鼓勵從繁瑣的網絡和緊湊型網絡生成的分割映射之間的對齊。
3. Approach
圖像語義分割是從 C 類預測圖像中每個像素的類別標簽的任務。分割網絡以 W xhx3 大小的 RGB 圖像 i 為輸入, 然后計算大小為 w0xh0xn 的要素圖 F, 其中 N 是通道的數量。最后, 應用分類器計算 f 尺寸為 W0xh0xc 的分割圖 Q, 并將其作為分割結果, 將其采樣到輸入圖像的空間大小 Wxh。
3.1. Structured Knowledge Distillation
我們應用知識蒸餾 [15] 策略將繁瑣分割網絡 T 的知識轉移到一個緊湊的分割網絡 S, 以更好地訓練緊湊型分割網絡。除了一個簡單的方案, 像素式蒸餾, 我們提出了兩個結構化的知識蒸餾方案, 成對蒸餾和整體蒸餾, 從繁瑣的網絡轉移到緊湊型的結構化知識網絡。該管道如圖2所示。
像素式蒸餾。我們將分割問題視為單獨像素標記問題的集合, 并直接使用知識蒸餾來對齊從緊湊型網絡生成的每個像素的類概率。我們采用了一種明顯的方法 [15]: 將繁瑣模型產生的類概率用作訓練緊湊型網絡的軟目標。損失函數給出如下,
（KL散度，又稱相對熵）
Pair-wise distillation.在提高空間標記連續性的對等馬爾可夫隨機場框架的啟發下, 我們提出在像素之間轉移對等關系, 特別是我們方法中的pair-wise similarities（對相似性）。
在我們的實現中, 兩個像素之間的相似性只是從特征 fi 和 fj 計算
整體蒸餾 我們對從繁瑣而緊湊的網絡中產生的分割映射之間的高階關系進行了對齊。分割映射的整體嵌入被計算為表示形式。
我們采用有條件生成對抗性學習 [29] 來制定整體蒸餾問題。緊湊型網絡被認為是以輸入 RGB 圖像 I 為條件的生成器, 預測分割圖 q 被視為假樣本。我們預計, q s與 Qt 相似, qt 是教師預測的分割圖, 盡可能被視為真實樣本。Wasserstein 距離 [13] 被用來評估真正的分布和假分布之間的區別, 這是寫如下,
其中 E [··] 是期望運算符, D (·) 是嵌入網絡, 在 GAN 中充當鑒別器, 將 Q 和 I 一起投影到一個整體嵌入分數中。梯度懲罰滿足了利普希茨的要求。
分割圖和條件 RGB 圖像串聯在一起, 嵌入網絡 D. d 的輸入是一個具有五個卷的完全卷積神經網絡。在最后三層之間插入兩個self-attention模塊, 以捕獲結構信息 [52, 57]。這樣的鑒別器能夠生成一個整體嵌入, 表示輸入圖像和分割映射的匹配程度。
3.2. Optimization
整個目標函數由傳統的多類交叉熵損失 mc (S) 組成, 具有像素化和結構化蒸餾項
其中lamda1和lamda2設置為10和 0:1, 使這些損失值范圍可比。我們最大限度地減少了與緊湊分割網絡 s 的參數有關的目標函數, 同時, 對于鑒別器 D 的參數將其最大化, 通過迭代以下兩個步驟實現了這一點:
訓練鑒別器 D. 訓練鑒別器相當于最小化 lho (S;D). D 的目的是給教師網的真實樣本一個高嵌入分數, 從學生網給假樣本的低嵌入分數。
·訓練緊湊的分割網絡 S。考慮到鑒別器網絡, 目標是最大限度地減少與緊湊分割網絡相關的多類交叉熵損耗和蒸餾損耗:
是 l ho 的一部分 (S;D) 在公式3中給出, 我們期望 S 在 D 的評價下獲得更高的分數。
4. Implementation Details
網絡結構。我們采用了最先進的分割架構 PSPNet [56], 其中包含 ResNet101 [14] 作為繁瑣的網絡 (教師) T。我們研究了最近的公共緊湊型網絡, 并使用了幾種不同的體系結構來驗證蒸餾框架的有效性。我們首先將 ResNet18 視為一個基本的學生網絡, 并對其進行消融研究。然后, 我們使用了一個開源 Mobiletv2plus [25], 它基于 ImageNet 數據集上預先訓練的 Mobiletv2 [37] 模型。我們還測試了 ESPNetC [28] 和 ESPNet [28] 的結構, 它們非常緊湊, 復雜度較低。
訓練設置。本文中的大多數分割網絡都是通過小批量隨機梯度下降 (sgd) 和40000迭代的權重衰減 (0:005) 進行訓練的。
我們隨機將圖像切割成512x512 作為訓練輸入。在訓練過程中應用正常的數據擴充方法, 如隨機縮放 (從0:5 到 2:1) 和隨機翻轉。除此之外, 我們還按照相應出版物 [28] 中的設置來重現 ESPNet 和 ESPNet-C 的結果, 并在我們的蒸餾框架下訓練緊湊型網絡。
