使用多尺度注意力進行語義分割

關注公眾號后后臺回復“多尺度語義分割”，下載打包好的論文和代碼。

有一項重要的技術，通常用于自動駕駛、醫(yī)學成像，甚至縮放虛擬背景：“語義分割。這是將圖像中的像素標記為屬于N類中的一個(N是任意數(shù)量的類)的過程，這些類可以是像汽車、道路、人或樹這樣的東西。就醫(yī)學圖像而言，類別對應于不同的器官或解剖結(jié)構(gòu)。

NVIDIA Research正在研究語義分割，因為它是一項廣泛適用的技術。我們還相信，改進語義分割的技術也可能有助于改進許多其他密集預測任務，如光流預測(預測物體的運動)，圖像超分辨率，等等。

我們開發(fā)出一種新方法的語義分割方法，在兩個共同的基準：Cityscapes和Mapillary Vistas上達到了SOTA的結(jié)果。。IOU是交并比，是描述語義分割預測精度的度量。

在Cityscapes中，這種方法在測試集上達到了85.4 IOU，考慮到這些分數(shù)之間的接近程度，這相對于其他方法來說是一個相當大的進步。

在Mapillary上，使用單個模型在驗證集上達到了61.1 IOU，相比于其他的使用了模型集成最優(yōu)結(jié)果是58.7。

預測結(jié)果

研究過程

為了開發(fā)這種新方法，我們考慮了圖像的哪些特定區(qū)域需要改進。圖2顯示了當前語義分割模型的兩種最大的失敗模式：細節(jié)錯誤和類混淆。

在這個例子中，存在兩個問題：細節(jié)和類混淆。

• 第一張圖片中郵箱的細節(jié)在2倍尺度的預測中得到了最好的分辨，但在0.5倍尺度下的分辨很差。
• 與中值分割相比，在0.5x尺度下對道路的粗預測要比在2x尺度下更好，在2x尺度下存在類混淆。

我們的解決方案在這兩個問題上的性能都能好得多，類混淆幾乎沒有發(fā)生，對細節(jié)的預測也更加平滑和一致。

在確定了這些錯誤模式之后，團隊試驗了許多不同的策略，包括不同的網(wǎng)絡主干(例如，WiderResnet-38、EfficientNet-B4、xcepase -71)，以及不同的分割解碼器(例如，DeeperLab)。我們決定采用HRNet作為網(wǎng)絡主干，RMI作為主要的損失函數(shù)。

HRNet已經(jīng)被證明非常適合計算機視覺任務，因為它保持了比以前的網(wǎng)絡WiderResnet38高2倍分辨率的表示。RMI損失提供了一種無需訴諸于條件隨機場之類的東西就能獲得結(jié)構(gòu)性損失的方法。HRNet和RMI損失都有助于解決細節(jié)和類混淆。

為了進一步解決主要的錯誤模式，我們創(chuàng)新了兩種方法：多尺度注意力和自動標記。

多尺度注意力

在計算機視覺模型中，通常采用多尺度推理的方法來獲得最佳的結(jié)果。多尺度圖像在網(wǎng)絡中運行，并將結(jié)果使用平均池化組合起來。

使用平均池化作為一個組合策略，將所有尺度視為同等重要。然而，精細的細節(jié)通常在較高的尺度上被最好地預測，大的物體在較低的尺度上被更好地預測，在較低的尺度上，網(wǎng)絡的感受野能夠更好地理解場景。

學習如何在像素級結(jié)合多尺度預測可以幫助解決這個問題。之前就有關于這一策略的研究，Chen等人的Attention to Scale是最接近的。在這個方法中，同時學習所有尺度的注意力。我們將其稱為顯式方法，如下圖所示。

受Chen方法的啟發(fā)，我們提出了一個多尺度的注意力模型，該模型也學會了預測一個密集的mask，從而將多尺度的預測結(jié)合在一起。但是在這個方法中，我們學習了一個相對的注意力mask，用于在一個尺度和下一個更高的尺度之間進行注意力，如圖4所示。我們將其稱為層次方法。

這種方法的主要好處如下：

• 理論訓練成本比Chen方法降低了約4x。
• 訓練只在成對的尺度上進行，推理是靈活的，可以在任意數(shù)量的尺度上進行。

圖5顯示了我們的方法的一些例子，以及已學習的注意力mask。對于左邊圖片中郵箱的細節(jié)，我們很少關注0.5x的預測，但是對2.0x尺度的預測非常關注。相反，對于右側(cè)圖像中非常大的道路/隔離帶區(qū)域，注意力機制學會最大程度地利用較低的尺度(0.5x)，以及更少地利用錯誤的2.0x預測。

自動標記

改進城市景觀語義分割結(jié)果的一種常用方法是利用大量的粗標記數(shù)據(jù)。這個數(shù)據(jù)大約是基線精標注數(shù)據(jù)的7倍。過去Cityscapes上的SOTA方法會使用粗標注標簽，或者使用粗標注的數(shù)據(jù)對網(wǎng)絡進行預訓練，或者將其與細標注數(shù)據(jù)混合使用。

然而，粗標注的標簽是一個挑戰(zhàn)，因為它們是有噪聲的和不精確的。ground truth粗標簽如圖6所示為“原始粗標簽”。

受最近工作的啟發(fā)，我們將自動標注作為一種方法，以產(chǎn)生更豐富的標簽，以填補ground truth粗標簽的標簽空白。我們生成的自動標簽顯示了比基線粗標簽更好的細節(jié)，如圖6所示。我們認為，通過填補長尾類的數(shù)據(jù)分布空白，這有助于泛化。

使用自動標記的樸素方法，例如使用來自教師網(wǎng)絡的多類概率來指導學生，將在磁盤空間上花費非常大的代價。為20,000張橫跨19個類的、分辨率都為1920×1080的粗圖像生成標簽大約需要2tb的存儲空間。這么大的代價最大的影響將是降低訓練成績。

我們使用硬閾值方法而不是軟閾值方法來將生成的標簽占用空間從2TB大大減少到600mb。在這個方法中，教師預測概率 > 0.5是有效的，較低概率的預測被視為“忽略”類。表4顯示了將粗數(shù)據(jù)添加到細數(shù)據(jù)和使用融合后的數(shù)據(jù)集訓練新學生的好處。

最后的細節(jié)

該模型使用PyTorch框架在4個DGX節(jié)點上對fp16張量核進行自動混合精度訓練。

論文：https://arxiv.org/abs/2005.10821

代碼：https://github.com/nvidia/semanic-segmentation