Kaggle車輛邊界識別第一名解決方案：使用預訓練權重輕鬆改進U-Net

選自arXiv

機器之心編譯

本文作者之一 Vladimir Iglovikov 曾取得 Kaggle Carvana Image Masking Challenge 第一名，本文介紹了他使用的方法：使用預訓練權重改進 U-Net，提升圖像分割的效果。

代碼地址（包含預訓練權重）：https://github.com/ternaus/TernausNet

隨著處理密集計算的計算機硬體的發展和平民化，研究者能夠處理擁有數百萬參數的模型。卷積神經網路在圖像分類、目標識別、場景理解等領域都取得了極大的成功。對幾乎所有的計算機視覺問題，基於 CNN 的方法都優於其他技術，在很多情況下甚至超越了人類專家。目前，幾乎所有的計算機視覺應用嘗試使用深度學習技術來改進傳統方法。它們影響到我們的日常生活，且這些技術的潛在應用場景似乎也很驚人。

可靠的圖像分割是計算機視覺領域的重要任務之一。該問題對醫療圖像領域極為重要，可以提高診斷能力，在場景理解領域中有助於創造安全的自動駕駛汽車。密集圖像分割本質上是把圖像分為有意義的區域，可看作是像素級別的分類任務。處理此類問題最直接（也緩慢）的方法是手動分割圖像。然而，這種方法極為耗時，而且人類數據管理員不可避免地會出錯、存在不一致問題。自動化該過程能提供儘可能快的、系統性的圖像分割。該過程需要達到一定的準確率，以在生產環境中使用。

在過去幾年中，研究者提出了不同的方法來解決該問題：創造一種 CNN，為一整張輸入圖像在單個前向傳遞中生成分割圖。其中，最成功的最優方法基於全卷積網路（FCN）[2]。其核心思想是將 CNN 中的全連接層替換成卷積層，成為強大的特徵提取器，直接輸出空間特徵圖，而不是全連接層輸出的分類分數。然後，上採樣這些圖，生成密集的逐像素輸出。該方法以端到端的方式訓練 CNN，分割任意大小的輸入圖像。此外，在 PASCAL VOC 等標準數據集上，該方法極大地改進了分割的準確率。

後來，這一方法進一步被改進為 U-Net 神經網路 [4]，U-Net 架構使用跳過連接（skip connection）將低層特徵圖與高層特徵圖結合起來，帶來精確的像素級定位。在上採樣部分，大量的特徵通道向更高的解析度層傳播上下文信息。在衛星圖像分析、醫療圖像分析等二值圖像分割競賽中，這種類型的網路架構已經證明了自己。

在此論文中，作者展示了如何使用預訓練權重輕鬆改進 U-Net 的性能。作者還將它應用到航空圖像標註數據集 [8] 中，該數據集包含多個城市的高清航空圖像。這些圖像的每個像素都被標註為「建築」或「非建築」類別。該架構另一個成功應用案例與初始化方案是 Kaggle Carvana 圖像分割競賽 [9]，本論文作者之一使用它作為解決方案的一部分，獲得了第一名。

II. 網路架構

通常，U-Net 架構包含一個收縮路徑來捕捉上下文信息，以及一個對稱的擴張路徑以進行精準的定位（見圖 1）。收縮路徑遵循典型的卷積網路架構，即交替卷積和池化運算，並逐步下採樣特徵圖，同時逐層增加特徵圖的數量。擴張路徑的每個階段由一個特徵圖上採樣和緊隨的卷積構成。

圖 1：編碼器-解碼器神經網路架構，亦稱為 U-Net，使用無全連接層的 VGG11 作為編碼器。每個藍色的矩形塊代表一個經過一系列變換的多通道特徵圖。矩形的長度代表相對的圖尺寸（像素級），其寬度和通道數量成正比。左邊編碼器部分的通道數逐漸增加，而右邊解碼器部分的通道數逐漸減少。頂部的箭頭表示每個編碼層的信息遷移，並傳輸至對應的解碼層。

因此，擴張分支可以提高輸出的解析度。為了進行定位，擴張路徑通過跳過連接將上採樣特徵和來自收縮路徑的高解析度特徵結合起來 [4]。模型的輸出是一個逐像素掩碼，展示了每個像素的類別。該架構被證明對有限數據的分割問題很有用，示例參見 [5]。

U-Net 可以從相對較小的訓練集中學習。多數情況下，圖像分割的數據集由至多幾千張圖像構成，因為手動標記掩碼是非常繁重的工作。通常 U-Net 以隨機初始化權重開始訓練。眾所周知，要使網路訓練避免過擬合，數據集應該足夠大，包含數百萬張圖像。在 ImageNet [10] 數據集上訓練的網路被廣泛地用於其它任務的網路初始化。用這種方法，可以讓網路非預訓練的幾層（有時僅僅是最後一層）利用數據集的特徵進行學習。

我們使用 VGG 族 [11] 中非常簡單的 CNN 作為 U-Net 網路的編碼器，它由 11 個順序層構成，稱為 VGG11，參見圖 2。VGG11 包含 7 個卷積層，每個緊隨著一個 ReLU 激活函數和 5 個最大池化運算，每個運算之後將特徵圖減半。所有卷積層都有 3x3 的卷積核，通道數如圖 2 所示。第一個卷積層有 64 個通道，然後網路加深，每個卷積層和最大池化運算之後通道數會加倍，直到通道數變為 512。在之後的卷積層中，通道數不變。

圖 2：VGG11 網路架構。在這張圖中，每個卷積層之後是 ReLU 激活函數。每個框中的數字表示對應特徵圖的通道數量。

為了構建解碼器，我們移除所有的全連接層，並用包含 512 個通道的單一卷積層來替換它們，該層是網路的瓶頸中間部分，分離編碼器與解碼器。為了構建解碼器，我們使用兩倍於特徵圖大小的轉置卷積層，同時把通道數量減少一半。轉置卷積的輸出接著被連接到解碼器對應部分的輸出。由此得到的特徵圖通過卷積運算來處理，以保持通道數量與對稱編碼器項相同。上採樣步驟被重複 5 次以配對 5 個最大池化層，如圖 1 所示。嚴格來說全連接層可以採用任何大小的輸入，但是由於我們有 5 個最大池化層，每個層下採樣圖像兩次，只有可被 32（2^5）整除的圖像可以用作當前網路實現的輸入。

圖 3：三種以不同權重初始化的 U-Net 模型的 Jaccard 指標隨訓練 epoch 的變化。藍線表示隨機初始化權重的模型，橙線表示編碼器以在 ImageNet 上預訓練的 VGG11 網路權重初始化的模型，綠線表示整個網路在 Carvana 數據集上預訓練的模型。

圖 4：綠色像素的二進位掩膜表示分類族群（建築）。圖 A 展示初始圖像和疊加的真實掩膜。圖 B 到圖 D 表示使用三種權重初始化方案並訓練了 100 個 epoch 後得到的預測結果。圖 B 表示隨機初始化權重的模型，圖 C 中的模型使用隨機初始化權重，編碼器以在 ImageNet 上預訓練的 VGG11 網路權重進行初始化，圖 D 中的模型使用在 Carvana 數據集上預訓練的權重。

在本論文中，通過使用微調（fine-tuning）技術初始化網路編碼器的權重，我們展示了如何提升 U-Net 的性能。這種神經網路被廣泛用於圖像分割任務，並在許多二值圖像分割、競賽中取得了當前最優結果。微調已廣泛用於圖像分類任務，但是就我們所知還未用於 U-Net 類型的架構。對於圖像分割問題，微調應該是更自然的選擇，因為收集大量數據集（尤其是醫療圖像）並進行很好地標註是很困難的。此外，預訓練網路可以大幅減少訓練時間，同時有助於防止過擬合。考慮到存在更多先進的預訓練編碼器比如 VGG16 [11] 或任何預訓練的 ResNet 網路，我們的方法還可進一步提升。有了這些改進的編碼器，解碼器可以像我們使用的一樣簡單。

論文：TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation

論文鏈接：https://arxiv.org/abs/1801.05746

摘要：像素級的圖像分割是計算機視覺中的艱巨任務。由編碼器和解碼器構成的經典 U-Net 架構經常用於分割醫學影像、衛星影像等。通常，神經網路使用在大型數據集（例如 ImageNet）上預訓練的網路權重進行初始化，相比用小型數據集從零開始訓練的網路能獲得更好的性能。在某些特定的應用中，特別是醫學和交通安全，模型的準確率至關重要。在本文中，我們展示了 U-Net 類型的架構如何利用預訓練的編碼器提升性能。我們的代碼和相關的預訓練權重已開源。我們比較了三種權重初始化方案：LeCun uniform、取自 VGG11 權重的編碼器和在 Carvana 數據集上訓練的完整網路。該網路架構是 Kaggle 競賽（Carvana Image Masking Challenge）中獲勝解決方案（在 735 名參賽者中排名第一）的一部分。

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

?------------------------------------------------

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！