文章學習29“Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning”

weixin_33670713發表於2018-10-08

這篇文章是CVPR2018年的作品,商湯科技出品。總得來說這篇文章在效能上並沒有什麼提示,但是把網路輕量化了,減少為原來1/3的引數量,還用了強化學習的策略。在這篇文章裡,作者認為真實的需要恢復的影象被進行了模糊,噪聲和JPEG壓縮三個步驟,所以在影象恢復時也考慮這三個處理過程並且將其分開,分成不同的小步驟依次進行處理。

作者在進行實驗時發現不同的小過程的處理級別和步驟會對最後的結果產生很大的影響,比如下圖,不同顏色是不同處理級別,相比於c來說,不同的處理方式、步驟、級別都會對最後的處理結果產生很大影響。所以在處理策略上需要下功夫。

8771353-3db21bd6a98a5aff.png

整個網路就由兩部分組成,一個是完成各種影象恢復作用的工具箱,另一個就是一個迴圈結構的agent,用以在每個階段選擇哪個工具。

8771353-6641948ee42118db.png

1)工具箱。本文提出的工具箱中包含了12個針對不同降質型別的CNN。每一種工具解決一種特定程度的高斯模糊、高斯噪聲、JPEG失真,這些失真在影象復原領域中最為常見。針對輕微程度失真的復原工具CNN僅有3層,而針對嚴重程度失真的工具達到8層。為了增強復原工具的魯棒性,本文在所有工具的訓練資料中均加入了輕微的高斯噪聲及JPEG失真。

8771353-308a71c64d61d897.png

2)準備好工具箱之後就是怎麼使用這些工具的演算法了,作者在這裡採用了強化學習的策略,將psnr的增值作為獎勵進行訓練,具體過程如上圖。作者花了很大的篇幅解釋了文章所用名詞的含義。比如“動作”,是指在每一個影象恢復步驟中採取的工具的這個過程,因為本文中有12(上表)個工具,所以每一個步驟就有13(12+停止)個動作;“回報”就是指的強化學習裡的reward,這裡採用的是psnr的增益。

網路結構就如上圖所示,在每一個步驟裡都需要有一個agent用來判斷應該使用toolbox裡的哪一個,agent由三部分組成:特徵提取器(Feature Extractor),包含了4個卷積層和1個全連線層,將輸入影象轉化為32維特徵向量;One-hot編碼器(One-hot Encoder),其輸入是前一步驟的動作估值向量,輸出將其轉換為對應的特徵向量;LSTM,其以前兩個模組輸出作為輸入,這個模組不僅觀測當前步驟的狀態特徵,還儲存了歷史狀態的資訊,該模組最後輸出當前步驟的估值向量,用於復原工具的選取。就這樣影象一步步的傳輸下去,經過幾個步驟之後遇到停止動作時即停止影象恢復。但在本文作者設定了一個最大傳輸步驟,即到達這個步驟之後即使沒有停止動作也停止訓練。

整個訓練過程是一個端到端的過程,演算法如下:

8771353-f5526d974ecfe185.png

對於每一張輸入影象,先通過所選取的工具鏈前向傳播得到最後的復原影象,通過與清晰參考影象對比得到MSE損失,然後通過工具鏈對誤差進行反向傳播,根據平均的梯度值更新工具網路的引數。

實驗:使用DIV2K訓練集的前750張影象用於訓練,後50張影象用於測試。通過摳取解析度為63x63的子影象,共得到25萬張訓練影象和3,584張測試影象。訓練時在每一張影象上隨機加上不同程度的高斯模糊、高斯噪聲和JPEG壓縮。演算法在訓練樣本中排除一些極度輕微或嚴重的失真,使用中度失真的影象進行訓練,而在輕度、中度和重度失真的影象上進行測試。測試結果如下:

8771353-c52348c10a51ee75.png

和VDSR和DnCNN在影象恢復上進行對比,引數少了2/3,但效能仍有提升。(VDSR-s指和RL-Restore同樣引數量的VDSR)

8771353-516eff54e5aed98e.png

在真實影象上的實驗結果如下圖所示,顯示本文所提出的方法比VDSR效果要好

8771353-9db4c2f599b9741a.png

相關文章