機器之心專欄
機器之心編輯部
在這篇論文中，研究者提出了一種新的、全監督語義分割訓練正規化「畫素對比學習」，強調利用訓練集中、跨影象的畫素 - 畫素對應關係來學習一個結構化的特徵空間，用來替代傳統的、基於影象的訓練正規化。
論文地址：https：//arxiv。org/pdf/2101。11939。pdf
程式碼地址：https：//github。com/tfzhou/ContrastiveSeg
當前，語義分割演算法的本質是透過深度神經網路將影象畫素對映到一個高度非線性的特徵空間。然而，現有演算法大多隻關注於區域性上下文資訊（單個影象內、畫素之間的位置和語義依賴性），卻忽略了訓練資料集的全域性上下文資訊（跨影象的、畫素之間的語義相關性），因而難以從整體的角度對習得的特徵空間進行約束，進而限制了語義分割模型的效能。
最近，蘇黎世聯邦理工學院及商湯研究院的研究者提出了一種新的、全監督語義分割訓練正規化：畫素對比學習（pixel-wise contrastive learning），強調利用訓練集中、跨影象的畫素 - 畫素對應關係（cross-image pixel-to-pixel relation）來學習一個結構化（well structured）的特徵空間，用來替代傳統的、基於影象的（image-wise）訓練正規化。
該訓練策略可直接應用於主流的語義分割模型，並在模型推理階段不引入額外計算開銷。下圖展示了在 Cityscapes 驗證集上主流分割演算法的效能，可以看出，在 DeepLabV3、HRNet、OCR 上引入畫素對比學習後，取得了較為顯著的效能提升。
當前語義分割領域忽略了什麼問題？
影象語義分割旨在為影象中的每個畫素預測一個語義標籤，是計算機視覺領域的一個核心問題。自 Fully Convolutional Network（FCN）［1］ 提出以後，主流的語義分割演算法強調影象內部的上下文資訊（intra-image context）。主要從兩方面入手： 1） 提出不同的上下文聚合模組（context aggreation module），如 dilated convolution， spatial pyramid pooling， encoder-decoder 及 non-local attention 等經典模型，核心思想是利用額外的模型引數或特殊的操作來建模提取影象內部的上下文資訊；2）傳統演算法將語義分割視為一個畫素級的分類任務，因而逐畫素獨立計算交叉熵損失（cross-entropy loss），但完全忽略了畫素間的依賴關係。因而有研究者提出結構化的（structure-aware）損失函式，如 pixel affinity loss ［2］，lovasz loss ［3］ 等，直接在訓練目標函式中對分割結果的整體結構資訊進行顯式約束。
然而，以上工作都只關注影象內部的上下文資訊，卻忽略了跨影象的、全域性的上下文資訊：在訓練集中，來自不同影象的畫素之間也具有極強的相關性，如圖 1（b），相同顏色的畫素表示他們具有相同的語義。
圖1：語義特徵空間學習示意圖。
更進一步，當前語義分割演算法的本質是透過深度神經網路，將影象畫素對映到一個高度非線性的特徵空間（如圖 1（c）），在這個過程中，只利用了上下文聚合模組或結構化的損失函式來強調區域性的畫素間的依賴關係，但忽略了一個本質問題：
一個理想的語義分割特徵空間究竟是什麼樣的？
研究者認為，一個好的分割特徵空間應同時具備兩個性質：
強判別能力：在該特徵空間中，每個畫素的特徵應具有較強的分類能力（strong categorization ability of individual pixel embeddings）；
高度結構化：同類畫素的特徵應高度緊緻（intra-class compactness），不同類畫素的特徵儘量分散（inter-class dispersion）。
然而，當前語義分割方法一般只關注性質 1，卻忽略了 2。此外，很多表徵學習（representation learning）的工作 ［4， 5］ 也驗證了：透過強調性質 2，有助於更好的增強性質 1。因此我們大膽假設，儘管當前的語義分割演算法已經取得了極佳效能，但是透過同時考慮性質 1 和 2，有可能習得一個更好的、結構化的分割特徵空間，進而進一步提高語義分割演算法的效能。
無監督對比學習引發的思考
近年來，無監督學習領域迎來了巨大發展，源頭是對比學習（contrastive learning）［6， 7］ 在海量無標註訓練樣本下的成功應用。假設 是一張沒有標註的訓練樣本影象 I 的特徵向量，為影象 I 的一個正樣本（positive sample）特徵，這個正樣本往往是對 I 施以某種變換得到的（如 flipping、cropping 操作等），為一個負樣本（negative sample）特徵，訓練集中其它非 I 的影象均被視為負樣本。之後透過對比學習損失函式，如下面的 InfoNCE loss ［8］ ，進行無監督訓練：
目標是將正樣本從大量負樣本中識別出來。無監督訓練得到的影象特徵表現出了極強的泛化能力，能夠為下游任務提供極佳的網路初始化權重，或僅經過少量有標籤樣本的微調（finetuning），即可獲得接近全監督訓練效能的影象分類模型。
無監督對比學習的成功也帶來了啟發，對比學習屬於度量學習（metric learning），本質是利用了資料集的整體資訊，學習了一個具有極強表達能力的影象表徵空間。在影象語義分割的全監督訓練條件下，訓練影象的每個畫素的標籤已經給出，我們可以將正樣本視為同屬於相同語義類別的畫素，將負樣本視為不屬於同一語義類別的畫素，而不論它們是否來源於同一個訓練影象。之後即可以利用度量學習或對比學習，對傳統的交叉熵損失加以改進，進而挖掘所有訓練影象中、畫素與畫素之間的全域性語義關係，進而獲得一個高度結構化的分割特徵空間，從而同時強調性質 1 和 2。由此研究者提出了一個基於畫素對比學習的、全監督的、語義分割訓練正規化，畫素對比學習（pixel-wise contrastive learning），強調利用訓練資料集的全域性上下文資訊，從整體的角度對習得的特徵空間進行顯式約束，使其在全域性結構上具備良好的性質（intra-class compactness 和 inter-class dispersion）。
如圖 1（d）所示，給定一個訓練樣本中的畫素 i，也稱之為錨點（anchor point），研究者將 i 與其它畫素在分割特徵空間上進行比較，儘可能拉近 i 與其同類的畫素（正樣本）的距離，並迫使 i 儘可能遠離其它不同類畫素（負樣本）。因此該訓練正規化能夠考慮整個訓練集中所有畫素的全域性語義相似度，使得模型能夠利用更具多樣性的、大規模樣本提升表徵學習的能力，從而獲得更優的語義特徵空間（如圖 1（e））。
基於 Pixel-Wise 交叉熵的經典語義分割損失函式有何問題？
下面以語義分割領域經典的 Pixel-wise 交叉熵作為出發點，進一步探討將度量學習或對比學習引入語義分割訓練的必要性。
如前所述，當前語義分割演算法將該任務視作一個逐畫素的語義分類問題，即為影象中的每個畫素 i 預測一個語義標籤 c。因此使用 pixel-wise 交叉熵作為訓練目標：
這裡 y 表示：透過 FCN 得到對畫素 i 的未歸一化的分類預測（unnormalized categorical score vector），也稱為 logit，表示畫素 i 的真實標籤，為的 one-hot 編碼。
然而，該最佳化目標函式有兩個缺點：
只對每個畫素的預測獨立地進行約束，而忽略了畫素之間的關係 ［9］ ；
由於使用了 softmax 操作，交叉熵的計算實際只依賴於 logits 之間的相對關係，卻無法直接約束習得的畫素特徵（cannot directly supervise on the learned representations） ［10］ 。
儘管最近的一些結構化損失函式（如 pixel affinity loss， lovasz loss 等）意識到了缺點 1，但也只考慮了同一個影象內部的畫素依賴關係，卻忽略了不同影象間畫素的語義一致性。而缺點 2，在語義分割領域則鮮有提及。
基於全監督、畫素 - 畫素對比學習的語義分割訓練正規化
圖2：基於全監督、畫素對比學習的語義分割模型。
本文提出的畫素對比學習（pixel-wise contrastive learning）能夠較好的解決交叉熵損失函式的兩個缺點。在訓練過程中，對於任意畫素（錨點）i， 其正樣本為其它與之同類的畫素，負樣本為其它與之不同類的畫素。值得注意的是，對於錨點 i 的正負樣本的選擇不限於同一張影象。對畫素 i，對比損失函式定義為：
這裡 表示畫素 i 的所有正樣本畫素的特徵，表示畫素 i 的所有負樣本畫素的特徵。由上式可以看出，透過畫素 - 畫素對比學習，研究者直接在語義分割的特徵空間上，將同屬於一個語義類別的畫素拉近，同時迫使不同語義類別的畫素互相遠離，從而同時強調了交叉熵損失的兩個缺點。
最終的語義分割損失函式定義為：
交叉熵損失促使分割模型學習具有判別力的特徵、提升分類能力（強調性質 1），pixel-wise contrastive loss 透過探索畫素之間的全域性語義關係，從整體上約束語義分割特徵空間（強調性質 2）。
下圖對只使用交叉熵損失（左圖）與上式混合損失函式（右圖）學習到的分割特徵進行了視覺化，可以看出，透過引入 pixel-wise contrastive loss，同類畫素特徵更加緊緻，而類與類之間可以更好的分離。這表明，透過結合一元（unary）交叉熵損失和二元（pari-wise）contrastive loss 的優勢，分割網路能夠學習到更好的特徵表示。
圖3：分割特徵視覺化。左圖：交叉熵損失；右圖：交叉熵 + 畫素對比損失。
基於以上框架，研究者進一步探索了畫素 - 區域對比學習及難例挖掘，並得出了若干有益結論，具體內容請移步原文。
實驗結果
為驗證所提方案的效能，研究者在三個標準的語義分割資料集（Cityscapes，PASCAL-Context，COCO-Stuff）上進行了實驗。以目前主流的語義分割模型（HRNet，OCR）為基礎，引入畫素對比學習後的效能對比如下：
實驗結果表明，在三個資料集上，本文提出的訓練正規化可以明顯地提升現有語義分割演算法的效能。此外，該演算法在模型部署階段，不會引入任何額外的計算開銷，不改變分割網路結構，具有較高的靈活性。
進一步探討
不同於當前主流演算法只關注影象內部畫素的區域性上下文資訊，本文提出跨影象、畫素對比損失函式，來挖掘訓練資料集中所有畫素的全域性關係，有效提升了語義分割的效能。這有助於我們重新思考當前主流的訓練正規化，不僅僅只關注訓練樣本自身的特性，同時從全域性的角度，關注訓練樣本之間關係。
本文也帶來了一些有益的啟示，如：
對比學習或度量學習依賴於正負樣本的質量，更智慧的取樣策略能夠幫助分割網路更快速有效的學習。
從度量學習的角度，交叉熵損失為一元損失函式（unary loss）， 而對比損失為二元損失函式（pair-wise loss），探索高階的度量損失函式有可能帶來更大的提升。
對比損失在計算中需要對正負樣本取樣，有可能借此更自然地實現訓練中的類別再均衡（class rebalance）。
本文方案在主流語義分割資料集上取得了有效的效能提升，並且有望在其它影象稠密預測任務中（如 2D 人體姿態估計，醫療影象分割等）發揮優勢。
參考文獻
［1］ Jonathan Long， Evan Shelhamer， and Trevor Darrell。 Fully convolutional networks for semantic segmentation。 In CVPR， 2015。
［2］ Tsung-Wei Ke， Jyh-Jing Hwang， Ziwei Liu， and Stella X Yu。 Adaptive affinity fields for semantic segmentation。 In ECCV， 2018
［3］ Maxim Berman， Amal Rannen Triki， and Matthew B Blaschko。 The lovasz-softmax loss： A tractable surrogate for the optimization of the intersection-over-union measure in neural networks。 In CVPR， 2018。
［4］ Weiyang Liu， Yandong Wen， Zhiding Yu， and Meng Yang。 Large-margin softmax loss for convolutional neural networks。 In ICML， 2016。
［5］ Florian Schroff， Dmitry Kalenichenko， and James Philbin。 Facenet： A unified embedding for face recognition and clustering。 In CVPR， 2015。
［6］ Alexey Dosovitskiy， Jost Tobias Springenberg， Martin Riedmiller， and Thomas Brox。 Discriminative unsupervised feature learning with convolutional neural networks。 In NeurIPS， 2014。
［7］ Ting Chen， Simon Kornblith， Mohammad Norouzi， and Geoffrey Hinton。 A simple framework for contrastive learning of visual representations。 ICML， 2020。
［8］ Michael Gutmann and Aapo Hyva rinen。 Noise-contrastive estimation： A new estimation principle for unnormalized statistical models。 In AISTATS， 2010。
［9］ Shuai Zhao， Yang Wang， Zheng Yang， and Deng Cai。 Region mutual information loss for semantic segmentation。 In NeurIPS， 2019。
［10］ Tianyu Pang， Kun Xu， Yinpeng Dong， Chao Du， Ning Chen， and Jun Zhu。 Rethinking softmax cross-entropy loss for adversarial robustness。 In ICLR， 2020。
AWS白皮書《策略手冊：資料、 分析與機器學習》
曾儲存過 GB 級業務資料
的組織現在發現，所儲存的資料量現已達 PB 級甚至 EB 級。要充分利用這 些海量資料的價值，就需要利用現代化雲資料基礎設施，從而將不同的資訊豎井融合統一。
無論您處於資料現代化改造過程中的哪個階段，本行動手冊都能幫助您完善策略，在整個企業範圍內高效擴充套件資料、分析和機器學習，從而加快創新並推動業務發展。
