當前位置:
首頁 > 知識 > TensorFlow實現StarGAN代碼全部開源,1天訓練完

TensorFlow實現StarGAN代碼全部開源,1天訓練完

來源:github、新智元

StarGAN 是去年 11 月由香港科技大學、新澤西大學和韓國大學等機構的研究人員提出的一個圖像風格遷移模型,是一種可以在同一個模型中進行多個圖像領域之間的風格轉換的對抗生成方法。近日,有研究人員將 StarGAN 在 TensorFlow 上實現的全部代碼開源,相關論文獲 CVPR 2018 Oral。

開源地址:https://github.com/taki0112/StarGAN-Tensorflow

作者:Junho Kim

看代碼之前,我們先來回顧一下 StarGAN 的原始論文。

StarGAN 對抗生成網路實現多領域圖像變換

圖像到圖像轉換(image-to-image translation)這個任務是指改變給定圖像的某一方面,例如,將人的面部表情從微笑改變為皺眉。在引入生成對抗網路(GAN)之後,這項任務有了顯著的改進,包括可以改變頭髮顏色,改變風景圖像的季節等等。

給定來自兩個不同領域的訓練數據,這些模型將學習如何將圖像從一個域轉換到另一個域。我們將屬性(attribute)定義為圖像中固有的有意義的特徵,例如頭髮顏色,性別或年齡等,並且將屬性值(attribute value)表示為屬性的一個特定值,例如頭髮顏色的屬性值可以是黑色 / 金色 / 棕色,性別的屬性值是男性 / 女性。我們進一步將域(domain)表示為共享相同屬性值的一組圖像。例如,女性的圖像可以代表一個 domain,男性的圖像代表另一個 domain。

一些圖像數據集帶有多個標籤屬性。例如,CelebA 數據集包含 40 個與頭髮顏色、性別和年齡等面部特徵相關的標籤,RaFD 數據集有 8 個面部表情標籤,如 「高興」、「憤怒」、「悲傷」 等。這些設置使我們能夠執行更有趣的任務,即多域圖像到圖像轉換(multi-domain image-to-image translation),即根據來自多個域的屬性改變圖像。

圖 1:通過從 RaFD 數據集學習遷移知識,應用到 CelebA 的多域圖像到圖像轉換結果。第一列和第六列顯示輸入圖像,其餘列是產生的 StarGAN 圖像。注意,圖像是由一個單一模型網路生成的,面部表情標籤如生氣、高興、恐懼是從 RaFD 學習的,而不是來自 CelebA。

在圖 1 中,前 5 列顯示了一個 CelebA 的圖像是如何根據 4 個域(「金髮」、「性別」、「年齡」 和 「白皮膚」)進行轉換。我們可以進一步擴展到訓練來自不同數據集的多個域,例如聯合訓練 CelebA 和 RaFD 圖像,使用在 RaFD 上訓練的特徵來改變 CelebA 圖像的面部表情,如圖 1 最右邊的列所示。

然而,現有模型在這種多域圖像轉換任務中既效率低,效果也不好。它們的低效性是因為在學習 k 個域之間的所有映射時,必須訓練 k(k-1)個生成器。圖 2 說明了如何訓練 12 個不同的生成器網路以在 4 個不同的域中轉換圖像。

圖 2: StarGAN 模型與其他跨域模型的比較。(a)為處理多個域,應該在每兩個域之間都建立跨域模型。(b)StarGAN 用單個生成器學習多域之間的映射。該圖表示連接多個域的拓撲圖。

為了解決這類問題,我們提出了StarGAN,這是一個能夠學習多個域之間映射的生成對抗網路。如圖 2(b) 所示,我們的模型接受多個域的訓練數據,僅使用一個生成器就可以學習所有可用域之間的映射。

這個想法很簡單。我們的模型不是學習固定的轉換(例如,將黑頭髮變成金色頭髮),而是將圖像和域信息作為輸入,學習將輸入的圖像靈活地轉換為相應的域。我們使用一個標籤來表示域信息。在訓練過程中,我們隨機生成一個目標域標籤,並訓練模型將輸入圖像轉換為目標域。這樣,我們可以控制域標籤並在測試階段將圖像轉換為任何想要的域。

我們還介紹了一種簡單但有效的方法,通過在域標籤中添加一個掩碼向量(mask vector)來實現不同數據集域之間的聯合訓練。我們提出的方法可以確保模型忽略未知的標籤,並關注特定數據集提供的標籤。這樣,我模型就可以很好地完成任務,比如利用從 RaFD 中學到的特徵合成 CelebA 圖像的面部表情,如圖 1 最右邊的列所示。據我們所知,這是第一個在不同的數據集上成功地完成多域圖像轉換的工作。

總結而言,這個研究的貢獻如下:

提出 StarGAN,這是一個新的生成對抗網路,只使用一個生成器和一個鑒別器來學習多個域之間的映射,能有效地利用所有域的圖像進行訓練。

演示了如何通過使用 mask vector 來學習多個數據集之間的多域圖像轉換,使 StarGAN 能夠控制所有可用的域標籤。

使用 StarGAN 在面部屬性轉換和面部表情合成任務提供了定性和定量的結果,優於 baseline 模型

圖 3:StarGAN 的概觀,包含兩個模塊:一個鑒別器 D 和一個生成器 G。(a)D 學習區分真實圖像和假圖像,並將真實圖像分類到相應的域。(b)G 接受圖像和目標域標籤作為輸入並生成假圖像。 (c)G 嘗試在給定原始域標籤的情況下,從假圖像中重建原始圖像。(d)G 嘗試生成與真實圖像非常像的假圖像,並通過 D 將其分類為目標域。

實驗結果

圖4:CelebA 數據集上面部屬性轉換的結果對凱勒巴數據集。第1列顯示輸入圖像,後4列顯示單個屬性轉換的結果,最右邊的列顯示多個屬性的轉換結果。H:頭髮的顏色;G:性別;A:年齡

圖5:RaFD 數據集上面部表情合成的結果

圖6:StarGAN-SNG 和 StarGAN-JNT 在 CelebA 數據集上的面部表情合成結果。

TensorFlow模型的實現

要求:

Tensorflow 1.8

Python 3.6

> python download.py celebA

下載數據集

> python download.py celebA

訓練

python main.py --phase train

測試

python main.py --phase test

celebA 測試圖像和你想要的圖像同時運行

預訓練模型

下載 celebA_checkpoint

結果 (128x128, wgan-gp)

女性

男性

預訓練權重:https://drive.google.com/open?id=1ezwtU1O_rxgNXgJaHcAynVX8KjMt0Ua-

訓練時間:少於 1 天

硬體:GTX 1080Ti

開源地址:https://github.com/taki0112/StarGAN-Tensorflow

- 加入AI學院學習 -


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

本站內容充實豐富,博大精深,小編精選每日熱門資訊,隨時更新,點擊「搶先收到最新資訊」瀏覽吧!


請您繼續閱讀更多來自 AI講堂 的精彩文章:

美國超算將重奪全球第一,性能超神威·太湖之光60%
谷歌發布AI 使用七原則,但仍與美國防部合作

TAG:AI講堂 |