首頁>科技>

使用CycleGAN和谷歌的QuickDraw資料集建立節日祥龍

今天是大年初一是2021年農曆新年的開始。在農曆新年裡,我們會經常看到五顏六色的舞龍。這裡我們使用ML的風格來慶祝我們的新年,我使用CycleGAN生成了帶有有趣裝飾光斑的龍(如上所示),我將在本文中介紹其實現步驟。也許這些由GAN產生的龍也會帶來好運!祝大家“恭喜發財”,在新的一年中幸福快樂。我永遠不會錯過使用俗氣雙關語的機會,所以GAN Hay Fat Choy!(粵語)

CycleGAN論文是由Berkeley AI Research的研究人員撰寫的,他們也釋出了程式碼,這使得該模型的實現非常容易實現。他們的模型可以使用週期一致的對抗網路實現不成對的影象到影象的轉換。在進行論文釋出之前工,在應用所需的影象轉換時,影象到影象的轉換需要成對影象的大型資料集。例如,要建立夏季到冬季的照片翻譯,您需要在冬季和夏季擁有完全相同影象的大型資料集。但是使用CycleGAN,您只需要夏季風景的資料集A和冬季風景的資料集B,該模型可以學習影象轉換,這是非常令人驚歎的創作!

實現CycleGAN

要測試CycleGAN,需要的是影象的資料集a和資料集B,以及一些計算能力。我這個專案的目標是創造節日的龍,用紅色、金色和綠色祥龍,就像這隻春節的龍。所以,我需要一個龍的資料集a和一個彩色的,響亮的耀斑的資料集B。

根據作者的說法,一個200到1000張圖片的資料集應該足夠讓CycleGAN訓練,儘管如果你不介意長時間的訓練,越多越好。對於龍族,我使用了谷歌的QuickDraw資料集,它包含了來自他們主辦的繪圖遊戲的眾包草圖。給定一個單詞,玩家必須在20秒內畫出它。他們已將所有產生的圖形作為開放資料集釋出。我是QuickDraw的忠實擁護者,因為它使我們能夠窺視世界上不同概念的集體視覺表示。時間限制迫使我們專注於僅包括在視覺上描述給定單詞的最重要的細節,從而留出時間來僅繪製概念的實質。翻看許多龍的圖畫,我始終看到翅膀,爬行動物的尾巴和噴火的氣息,我想這是龍的本質。在短短20秒內完成的一些龍形畫使我印象深刻。對於我的資料,我為資料集A收集了大約450個龍素描。然後,對於資料集B,我建立了一個包含450個隨機形狀和填充的資料集-如果願意,可以使用隨機光斑。我在下圖中說明了CycleGAN如何使用我的資料集來高階學習影象轉換

克隆作者的Github之後,需要在資料集下建立一個資料夾。 在該資料夾中,建立“ trainA”,“ trainB”,“ testA”和“ testB”資料夾以容納您的資料。 資料放置到位後,您可以執行以下程式碼行,根據需要修改資料集位置和模型名稱,以開始訓練模型.

python train.py --dataroot ./datasets/sketch2festivedragon --name cycle_gan_dragon_rg --model cycle_gan

在NVIDIA GeForce RTX 3090上,批處理大小為1(預設設定),該模型使用約10GB的GPU記憶體,並在4個多小時內完成了訓練(200個輪次)。 如果您使用作者的程式碼,我強烈建議您遵循他們的訓練技巧,並執行以下命令以在本地圖形Web伺服器上啟動實時訓練視覺化。

python -m visdom.server

眾所周知,訓練GAN十分困難,因為損失圖很難解釋,也不一定與期望的結果保持一致。 相反,將結果視覺化非常重要,以便跟蹤實現所需影象轉換的進度。 如果您經常進行測試,您甚至可能會在幾個不同的時期看到許多效果不錯的例項。 預設情況下,程式碼每5個時間段儲存一次模型的權重。 訓練後,您可以使用任何這些檢查點權重返回並測試模型的輸出。 您所要做的就是執行以下程式碼行,並修改資料集的位置,模型的名稱和輪次。 測試影象將儲存在結果資料夾中。

python test.py --dataroot ./datasets/sketch2festivedragon --name cycle_gan_dragon_rg --model cycle_gan --epoch 35

觀察每輪次的測試輸出,我發現在整個訓練過程中都有可愛的裝飾龍。 在我尋找藝術耀斑的時候,龍與耀斑資料集完全相似並不重要。 訓練快要結束時,龍開始顯得太扭曲了,這在模型試圖扭曲輸入資料以匹配耀斑資料集時是有道理的。 我整理了一些收藏,將它們顯示在本文開頭的影象中,經過幾輪訓練時表現得不同效果。 我特別喜歡在第35輪次前後生成的龍,其中模型的變形為龍的草圖增加了熾烈而細膩的色彩飛濺。 以下是35輪的一些示例。

網路對如何將隨機的彩色形狀扭曲為輸入影象以形成類似龍的圖形的印象也給我留下了深刻的印象。 以下是訓練結束時網路繪製的一些龍頭樣本。

今年是牛年,因此我們也要生成一些喜慶的牛,。 我收集了約170張類似牛的QuickDraw牛圖,並重復了上述步驟,以便在這個新的迷你資料集上訓練CycleGAN。 訓練在大約一個小時40分鐘內完成。 以下是一些我最喜歡的示例輸出。

只是為了好玩,我們龍資料集上訓練的模型在牛繪圖輸入上的表現。 以下是使用35輪龍和牛的示例-噴火+

9
最新評論
  • 整治雙十一購物亂象,國家再次出手!該跟這些套路說再見了
  • 從產品使用角度,談談客服機器人帶來的溫暖