Kaggle競賽第一名解決方案：使用預訓練權重輕鬆改進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 名參賽者中排名第一）的一部分。