CVPR 2018 | 商湯科技Spotlight論文詳解：RL-Restore普適圖像復原演算法

在底層視覺演算法領域，商湯科技發表了一篇亮點報告（Spotlight）論文，提出基於深度強化學習的RL-Restore演算法，解決具有複雜混合失真的圖像復原問題。

論文：Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning

作者：Ke Yu, Chao Dong, Liang Lin, Chen Change Loy

論文鏈接：

https://arxiv.org/abs/1804.03312

Project Link：

http://mmlab.ie.cuhk.edu.hk/projects/RL-Restore/

簡介

在底層視覺演算法領域，卷積神經網路（CNN）近年取得了巨大的進展，在諸如去模糊、去噪、去JPEG失真、超解析度等圖像復原任務上已經達到了優異的性能。但是現實圖像中的失真往往更加複雜，例如，經過多個圖像降質過程後，圖像可能包含模糊、雜訊、JPEG壓縮的混合失真。這類混合失真圖像對目前的圖像復原演算法仍然充滿挑戰性。

近期的一些圖像復原工作（如VDSR、DnCNN等）證實了一個CNN網路可以處理多種失真類型或不同失真程度的降質圖像，這為解決混合失真問題提供了新的思路。但是，這類演算法均選用了複雜度較高的網路模型，帶來了較大的計算開銷。另外，這些演算法的網路均使用同一結構處理所有圖像，未考慮一些降質程度較低的圖像可以使用更小的網路進行復原。

針對現有圖像復原CNN演算法模型複雜，計算複雜度高的問題，本文提出的RL-Restore演算法彌補了這些不足，以更加高效靈活的方式解決了複雜的圖像復原問題。

RL-Restore演算法的設計思想與挑戰

當前流行的圖像復原理念認為解決複雜的圖像復原問題需要一個大型的CNN，而本文提出了一種全新的解決方案，即使用多個小型CNN專家以協作的方式解決困難的真實圖像復原任務。RL-Restore演算法的主要思路是設計一系列基於小型CNN的復原工具，並根據訓練數據學習如何恰當地組合使用它們。這是因為現實圖像或多或少受到多種失真的影響，針對複雜失真的圖像學習混合使用不同的小型CNN能夠有效的解決現實圖像的復原問題。不僅如此，因為該演算法可以根據不同的失真程度選取不同大小的工具，相較於現有CNN模型，這一新方法使用的參數更少，計算複雜度更低。

RL-Restore演算法的目標是對一張失真圖像有針對性地選擇一個工具鏈（即一系列小型CNN工具）進行復原，因而其該演算法包含了兩個基本組件：

• 一個包含多種圖像復原小型CNN的工具箱；
• 一個可以在每一步決定使用何種復原工具的強化學習演算法。

本文提出的工具箱中包含了12個針對不同降質類型的CNN（如表1所示）。每一種工具解決一種特定程度的高斯模糊、高斯雜訊、JPEG失真，這些失真在圖像復原領域中最為常見。針對輕微程度失真的復原工具CNN僅有3層，而針對嚴重程度失真的工具達到8層。為了增強復原工具的魯棒性，本文在所有工具的訓練數據中均加入了輕微的高斯雜訊及JPEG失真。

有了工具箱，如何選擇工具成為本文解決的主要挑戰之一。圖1展示了不同工具鏈的圖像復原結果，可以看到對工具鏈的微小調整可能導致復原結果的劇烈變化。本文解決的第二個挑戰在於，沒有一個已有的工具可以恰當的處理「中間結果」。例如，去模糊的工具可能也會放大雜訊，導致後面已有的去噪工具無法有效處理新引入的未知失真。針對這些挑戰，本文使用強化學習演算法訓練得到有效的工具選擇策略，同時還提出聯合訓練演算法對所有工具進行端到端的訓練以解決有效復原「中間結果」的挑戰。

基於強化學習的普適圖像復原

RL-Restore演算法的框架（如圖2所示）。對於一張輸入圖像，agent首先從工具箱中選擇一個工具對它進行恢復。然後agent根據當前的狀態（包括復原中間結果和之前工具的選擇）來選取下一個動作（使用復原工具或停止），直到演算法決定終止復原過程。

動作（action）：在每一個復原步驟 t，演算法會輸出一個估值向量vt選擇一個動作at。除了停止動作以外，其餘每一個動作均代表使用某個復原工具。在本文中，工具箱內共包含12個工具，因而演算法總共包含13個動作。

狀態（state）：狀態是演算法可以觀測到的信息，在步驟t的狀態記為St={It,v ?t}，其中It是當前步驟的輸入圖像，v ?t=vt-1是前一步驟的動作估值向量，包含了前一步驟的決策信息。

回報（reward）：在強化學習中，演算法的學習目標是最大化所有步驟的累積回報，因而回報是驅動演算法學習的關鍵。本文希望確保圖像質量在每一步驟都得到提升，因此設計了一個逐步的回報函數rt=Pt+1-Pt，其中Pt+1和Pt分別代表步驟t的輸入圖像和輸出圖像的PSNR，度量每個步驟中圖像PSNR的變化。

結構：虛線框內的agent包含了三個模塊（如圖2所示）：

• 特徵提取器（Feature Extractor），包含了4個卷積層和1個全連接層，將輸入圖像轉化為32維特徵向量；
• One-hot編碼器（One-hot Encoder），其輸入是前一步驟的動作估值向量，輸出將其轉換為對應的特徵向量；
• LSTM，其以前兩個模塊輸出作為輸入，這個模塊不僅觀測當前步驟的狀態特徵，還存儲了歷史狀態的信息，該模塊最後輸出當前步驟的估值向量，用於復原工具的選取。

訓練：每一個復原工具的訓練均使用MSE損失函數，而agent的訓練則使用deep Q-learning演算法。由於LSTM具有記憶性，每一個訓練樣本均包含一條完整的工具鏈。

聯合訓練演算法

至此，RL-Restore演算法已經擁有了較好的工具選取策略，還需要解決對「中間結果」進行復原的挑戰。前文已經提到，由於前面的復原步驟可能引入新的未知失真，沒有一個已有工具能對這類複雜的「中間結果」進行有效處理。因此，本文提出了聯合訓練演算法，將所有的工具以及工具的選擇進行端到端地訓練，從而解決「中間結果」的復原問題。具體而言，對於每一張輸入圖像，先通過所選取的工具鏈前向傳播得到最後的復原圖像，通過與清晰參考圖像對比得到MSE損失，然後通過工具鏈對誤差進行反向傳播，根據平均的梯度值更新工具網路的參數。

實驗結果

本文使用DIV2K訓練集的前750張圖像用於訓練，後50張圖像用於測試。通過摳取解析度為63x63的子圖像，共得到25萬張訓練圖像和3,584張測試圖像。本文在每一張圖像上隨機加上不同程度的高斯模糊、高斯雜訊和JPEG壓縮。演算法在訓練樣本中排除一些極度輕微或嚴重的失真，使用中度失真的圖像進行訓練（如圖3所示），而在輕度、中度和重度失真的圖像上進行測試。

本文與現有的VDSR和DnCNN圖像復原演算法相比，模型複雜度更低而復原性能更加優異（如表2、3所示）。其中VDSR-s是與VDSR結構相似的小參數模型，其參數量與RL-Restore演算法相當。表2展示了RL-Restore演算法具有最小的參數量和計算複雜度，表3展示了RL-Restore演算法與VDSR和DnCNN等大模型在輕度和中度失真測試集上具有類似的性能，而在重度失真測試集上則表現得更加優異。在參數量相當的情況下，RL-Restore演算法在各個測試集上均比VDSR-s演算法擁有更加優異的復原性能。圖4展示了不同演算法和本文演算法在不同步驟復原結果的對比。

本文也使用實際場景圖像對RL-Restore演算法進行了進一步測試。如圖5所示，測試圖像由智能手機採集，其中包含了模糊、雜訊和壓縮等失真，直接使用訓練好的RL-Restore和VDSR模型在這些真實場景圖像進行測試。由結果可以看到，RL-Restore演算法取得了明顯更加優異的復原結果，圖5(a, c) 展示了RL-Restore演算法成功修復由曝光雜訊和壓縮帶來的嚴重失真；圖5(b, d, e) 展示了本文方法可以有效地處理混合的模糊與雜訊。

結論

本文提出了一種新穎的基於強化學習的圖像復原演算法—RL-Restore。與現有的深度學習方法不同，RL-Restore演算法通過學習動態地選取工具鏈從而對帶有複雜混合失真的圖像進行高效的逐步復原。基於合成數據與現實數據的大量實驗結果證實了該演算法的有效性和魯棒性。由於演算法框架的靈活性，通過設計不同的工具箱和回報函數，RL-Restore演算法為解決其他富有挑戰性的底層視覺問題也提供了新穎的解決思路。