SRGAN With WGAN:讓超分辨率算法訓練更穩定

 2018-05-21 15:01:00.0

作者丨胡志豪

單位 | 棒谷網絡科技圖像工程師

經歷丨實現過電商大規模以圖搜圖、超分辨率等項目

寫在前面

此文挖坑了很久,項目本在 7 個月前已經結束,但一直沒心思把代碼整理出來,後來發現有相同思路的團隊把它寫成論文佔坑了,也就更沒動力寫了。昨晚發現我的 Github 項目竟然有星星,感受到了莫大的支持,忽然燃起了寫文章的動力,於是就有了現在這篇。

SRGAN

SRGAN,2017 年 CVPR 中備受矚目的超分辨率論文,把超分辨率的效果帶到了一個新的高度,而 2017 年超分大賽 NTIRE 的冠軍 EDSR 也是基於 SRGAN 的變體。

SRGAN 是基於 GAN 方法進行訓練的,有一個生成器和一個判別器,判別器的主體使用 VGG19,生成器是一連串的 Residual block 連接,同時在模型後部也加入了 subpixel 模塊,借鑑了 Shi et al 的 Subpixel Network [6] 的思想,讓圖片在最後面的網絡層才增加分辨率,提升分辨率的同時減少計算資源消耗。詳細的介紹建議大家還是直接看論文 [1],網上也有一些解讀的文章,這裏就直接介紹實現的一些創新細節。

GAN存在的問題

傳統 GAN 存在一個問題是,你無法知道什麼時候該停止訓練 Generator,什麼時候該停止訓練 Discriminator,如果過度訓練 Discriminator,Generator 就無法學習下去,反之也會導致模型效果差。

如果能有一個 loss 指標來反映訓練情況的話,訓練的難度就會大大降低。而 17 年提出的WGAN [3] 正是解決這一問題的重要方法。 

WGAN 使用 Wasserstein 距離來描述兩個數據集分佈之間的差異程度,只要把模型修改成 WGAN 的形式,就能根據一個唯一的 loss 來監控模型訓練的程度。有關 WGAN 的解釋強烈推薦大家閱讀這篇文章:令人拍案叫絕的 Wasserstein GAN [4],作者用非常直白明瞭的語言介紹 WGAN

SRGAN結合WGAN

SRGAN 的一個超讚復現來自 @董豪 之手,他也是 tensorlayer 的作者之一,他的復現項目在 Github 上得到了大量的星星,而筆者的代碼就正是在他的基礎上進行拓展,首先非常感謝作者的開源。

· 判別器最後一層去掉 sigmoid 

· 生成器和判別器的 loss 不取 log 

· 每次更新判別器的參數之後把它們的絕對值截斷到不超過一個固定常數 c 

· 不要用基於動量的優化算法(包括 momentum 和 Adam),推薦 RMSProp,SGD 也行 

--來自《令人拍案叫絕的Wasserstein GAN

根據這篇文章的介紹,對代碼進行以上四項修改,把 GAN 的訓練方式轉換成 WGAN,可以在 tensorboard 中監控 loss 的下降情況。因此,筆者對項目原作的進行了一些修改: 

1. 對模型代碼進行了上面 WGAN 的改造; 

2. 增加了 Tensorboard,監控 loss 下降情況; 

3. 對作者的 model.py 中,Generator 的最後一層的卷積 kernel 從 1×1 改成 9×9,這是原論文建議的結構。 

SRGAN With Wasserstein GAN 的完整代碼已經開源到筆者的 Github [5],裏面有詳細的改動介紹和使用說明,覺得好的話就給個 Star 呀!

▲ 修改後的訓練Loss下降情況


下面是一些復現後的超分復原效果:

一個來自工業界的問題

在實際生產使用中,遇到的低分辨率圖片並不一定都是 PNG 格式的(無損壓縮的圖片復原效果最好),而且會帶有不同程度的失真(有損壓縮導致的 artifacts)。筆者嘗試過很多算法,例如 SRGAN、EDSR、RAISR、Fast Neural Style 等等,這類圖片目前使用任何一種超分算法都沒法在提高分辨率的同時消除失真

這個問題我在 @董豪 SRGAN 項目的 issue 中也討論過,同時在知乎也提出過這個問題:SRGAN 超分辨率方法對於低清 jpg 格式的圖片復原效果是否比不上對低清 png 格式的復原效果?

可惜沒有很好的答案。目前學術界貌似還沒有很好的算法,這裏歡迎各位在評論區或者 Github 上來討論。

文章來源:機器之心