Bye Bye TPU，4個(gè)GPU訓(xùn)練“史上最強(qiáng)”GAN！作者開源PyTorch模型

夏乙 安妮 2019-03-23 13:20:07 來源：量子位

安妮 夏乙 假裝發(fā)自 凹非寺

量子位 出品 | 公眾號 QbitAI

“不，你不想?！?/p>

每次被BigGAN史上最強(qiáng)”的效果吸引，想要用其他數(shù)據(jù)集訓(xùn)練一番，腦海深處都會響起這樣一個(gè)聲音。

就仿佛DeepMind團(tuán)隊(duì)訓(xùn)練BigGAN用的512個(gè)TPU，齊刷刷發(fā)出不懷好意的嘲笑。

現(xiàn)在，

好消息來了！

BigGAN一作Andrew Brock新鮮放出了只需4-8個(gè)GPU即可訓(xùn)練的新版模型，這一次，他用的是PyTorch。

雖然“只需”令人嫉妒，但比起原來動輒128個(gè)、512個(gè)TPU的硬件需求，簡直就扶貧濟(jì)困，平民之光。

網(wǎng)友們一夜之間把Brock這條Twitter轉(zhuǎn)了100多次，在Reddit論壇上也夸這個(gè)新實(shí)現(xiàn)”fantastic”、”awesome”、”cool”，表示一定要玩。

更讓人期待的是，打破了128個(gè)TPU的門檻之后，用其他數(shù)據(jù)集訓(xùn)練的BigGAN一定會層出不窮啊！風(fēng)景、waifu指日可待。

另外，吸引網(wǎng)友們的不僅僅是算力需求的降低，還有：

等你的PyTorch版實(shí)現(xiàn)等了好久了！

擺脫算力束縛

Brock開源的低成本版PyTorch實(shí)現(xiàn)，內(nèi)容包括訓(xùn)練、測試、采樣所需的腳本、把THHub原版BigGAN生成器權(quán)重遷移到PyTorch的腳本，還有完整的預(yù)訓(xùn)練檢查點(diǎn)。

除了現(xiàn)有的ImageNet預(yù)訓(xùn)練模型之外，他還打算放出Places-365數(shù)據(jù)集上預(yù)訓(xùn)練的模型。

如果你想用自己的數(shù)據(jù)來訓(xùn)練BigGAN，可以選擇從頭開始訓(xùn)練，也可以在ImageNet預(yù)訓(xùn)練模型的基礎(chǔ)上微調(diào)。無論如何，只需4-8個(gè)GPU。

Brock還在Reddit論壇上說，其實(shí)一個(gè)GPU也行，但就是要花的時(shí)間太長了。

原版是128-512個(gè)TPU，新版是4-8個(gè)GPU，這之間有巨大的算力差距，Brock到底是怎樣用Pytorch完整復(fù)現(xiàn)BigGAN的？

大思路就是：靠梯度累加，在小硬件上使用大批量（batch）。Brock已經(jīng)將整個(gè)復(fù)現(xiàn)過程詳細(xì)記錄了下來。

我們先來看看這個(gè)低成本版BigGAN怎么用。

開始之前，需要準(zhǔn)備三樣訓(xùn)練工具：

• PyTorch（版本1.0.1）
• tqdm，numpy，scipy和h5py
• ImageNet訓(xùn)練集

首先，你可以選擇為你的目標(biāo)數(shù)據(jù)集準(zhǔn)備一個(gè)預(yù)處理HDF5版本，以實(shí)現(xiàn)更快的I/O，這需要通過修改和運(yùn)行計(jì)算FID需要的Inception moment。運(yùn)行這行代碼就可以：

sh scripts/utils/prepare_data.sh

默認(rèn)情況下，它會假設(shè)你的ImageNet訓(xùn)練集下載到根目錄的data位置下，并將以128×128像素的分辨率準(zhǔn)備緩存的HDF5。

在腳本文件夾中，有多個(gè)bash腳本可以用來訓(xùn)練不同批大小的BigGAN。這些代碼默認(rèn)你無法訪問完整的TPU pod，因此，會運(yùn)用梯度累加技術(shù)來假冒大批量。

具體來說，先在多個(gè)小批量上將梯度平均，然后在N次累計(jì)后，再執(zhí)行優(yōu)化程序步驟。

在默認(rèn)情況下，launch_BigGAN_bs256x8.sh腳本訓(xùn)練一個(gè)全尺寸96通道的BigGAN模型，用批大小為256的8個(gè)梯度累加，得到總體大小為2048的批量。

用的硬件少了，難免要多花點(diǎn)時(shí)間。按上面的默認(rèn)情況，在128×128像素的ImageNet上訓(xùn)練BigGAN，用8個(gè)Tesla V100進(jìn)行全精度的無張量訓(xùn)練，大約需要15天完成150000次迭代。

正確的流程是，你首先需要做的就是設(shè)置可以支持的最大的批大小，確定后再修改腳本，使批大小×梯度累加的數(shù)量=所需的批大小總量（BigGAN默認(rèn)為2048）。

需要注意，梯度累加意味著需要將預(yù)估的奇異值（singular value，SV）和批歸一化（Batch Normalization，BN）統(tǒng)計(jì)量提高8倍。

這意味著，BN的統(tǒng)計(jì)值更接近于靜止?fàn)顟B(tài)，相較之下奇異值的估計(jì)值往往更加準(zhǔn)確。

因此，Brock在默認(rèn)情況下用測試模式下的G來衡量標(biāo)準(zhǔn)，即用BN的運(yùn)行狀態(tài)預(yù)估值來替代去計(jì)算靜態(tài)統(tǒng)計(jì)值。不過，最終repo也支持靜態(tài)統(tǒng)計(jì)，雖然可能會導(dǎo)致前期的梯度累加失效，但問題也不大。

標(biāo)準(zhǔn)和示例

在訓(xùn)練期間，這個(gè)腳本會輸出具有訓(xùn)練標(biāo)準(zhǔn)和測試標(biāo)準(zhǔn)的日志，并且將保存模型權(quán)重/優(yōu)化的程序參數(shù)的多個(gè)副本，并且每次保存權(quán)重時(shí)，都將生成示例和插值。

日志文件夾中存儲了處理這些日志的腳本，并將結(jié)果用Matlab繪制出來。

訓(xùn)練過程完成后，你可以用sample.py代碼塊生成其他示例和插值，還能改變截?cái)嘀?、批大小、靜態(tài)統(tǒng)計(jì)累計(jì)值等進(jìn)行測試。

默認(rèn)情況下，所有內(nèi)容都保存到權(quán)重/示例/日志/數(shù)據(jù)文件夾中，repo被假定到與它們在同一個(gè)文件夾里了。你可以用—base_root代碼塊將這些參數(shù)指向不同的基礎(chǔ)文件夾，或者用各自的參數(shù)選擇每個(gè)基礎(chǔ)文件夾的特定位置。

在Brock小哥給出的實(shí)現(xiàn)中，包含了BigGAN-deep的運(yùn)行腳本，但沒有用它完全訓(xùn)練成一個(gè)模型，是未經(jīng)測試的一個(gè)腳本。

此外，Brock還給出了包括在CIFAR上運(yùn)行的腳本，以及在ImageNet上運(yùn)行SA-GAN和SN-GAN的腳本。SA-GAN是假設(shè)用4個(gè)TitanX訓(xùn)練時(shí)的腳本，是在批大小為128時(shí)加2個(gè)梯度累加的情況下進(jìn)行的。

用自己的數(shù)據(jù)微調(diào)預(yù)訓(xùn)練模型

我們前邊也提到過，除了從頭開始訓(xùn)練，還可以拿Brock預(yù)訓(xùn)練的模型來微調(diào)。

要運(yùn)行的腳本和從頭訓(xùn)練一樣，但是要加上—resume。

繼續(xù)進(jìn)行中斷了的訓(xùn)練也是用這個(gè)腳本。

要準(zhǔn)備微調(diào)用的數(shù)據(jù)集，需要添加到datasets.py中，并在utils.py里調(diào)整convenience dicts，例如dset_dict、imsize_dict、root_dict、nclass_dict、classes_per_sheet_dict等等設(shè)置，這樣才能為數(shù)據(jù)集創(chuàng)建合適的元數(shù)據(jù)。

然后，重復(fù)prepar_data.sh的流程。

我有一塊1080Ti，行嗎？

低成本版一出，引來不少網(wǎng)友歡呼，很多一看就來自學(xué)神們。

他們很是關(guān)注Brock所用的方法，比如同步的批歸一化(synchronized batchnorm)、梯度累加(gradient accumulation)技術(shù)等等，還打算用到自己的研究里。

當(dāng)然，更多人關(guān)心的是：

用這個(gè)“官方非官方”PyTorch版實(shí)現(xiàn)訓(xùn)練，成本究竟怎么樣？

按照Brock舉的例子計(jì)算，8塊Tesla V100訓(xùn)練15天，在亞馬遜云服務(wù)AWS上的價(jià)格是8800多美元。

也就是說，如果你自己本地沒有那么多GPU，還是挺貴……

于是，評論區(qū)除了歡呼的大神們之外，也冒出了很多貧窮的聲音：

哇哦好贊??！但是你說得好像4-8塊GPU多便宜似的哈哈哈哈。

然而我只有一塊1080ti……

在Colab里哭。（量子位注：就是蹭免費(fèi)GPU的Colab啦）

也有人真誠提問：

我有一塊2080ti，你估計(jì)半個(gè)星期能訓(xùn)練到收斂嗎？還是需要幾周？

Brock給出的答案不容樂觀：一塊2080ti啊，那得好幾個(gè)月了。

也就是說，如果預(yù)算不充沛，還是拿預(yù)訓(xùn)練模型微調(diào)的好。

不過，評論區(qū)也有高人，為新實(shí)現(xiàn)找了個(gè)絕佳的使用場景：

那些之前買顯卡挖礦的人，終于有比較酷的事情可干了。

這一提議頓時(shí)有人響應(yīng)：

哈……我有12塊英偉達(dá)1080，就是因?yàn)檫@個(gè)。

如果你也有挖礦遺留下來的GPU，那就一切好辦了。

轉(zhuǎn)投PyTorch體驗(yàn)怎么樣？

除了省算力，“終于有了完整PyTorch版實(shí)現(xiàn)”也是網(wǎng)友們非常關(guān)心的。

畢竟，Brock作為一名（曾經(jīng)的）DeepMind實(shí)習(xí)生，改用PyTorch實(shí)屬難得。這一版本“官方非官方”實(shí)現(xiàn)是他和Alex Andonian一起完成的，而并沒有提到BigGAN另外兩位作者。

△?BigGAN一作Andrew Brock

他使用PyTorch體驗(yàn)怎樣？如何看待在PyTorch和TensorFlow中實(shí)現(xiàn)BigGAN的差異？

在Reddit上，也有網(wǎng)友向Brock好奇地提出了這個(gè)問題。

Brock坦言，用PyTorch訓(xùn)練BigGAN，是種非常可愛的體驗(yàn)。

他認(rèn)為PyTorch讓很多事情都變得容易多了，比如組成層、測試、調(diào)試、實(shí)現(xiàn)模型并行性等操作，一切都可以信手拈來。

整個(gè)實(shí)現(xiàn)的過程中，當(dāng)然也遇到了一些沒想到的困難，不過Brock說都是他自己的錯。比如直到開始檢查梯度時(shí)，才發(fā)現(xiàn)錯位的ReLU已經(jīng)造成了大面積破壞。

此外，他遇到的最大挑戰(zhàn)其實(shí)是算力不足：因?yàn)锽rock只有多所大學(xué)的共享服務(wù)器可用，因此調(diào)試/測試周期也拉長了不少。每次訓(xùn)練只能運(yùn)行24小時(shí)，等待下一次開始的時(shí)間，則不止24小時(shí)了。

傳送門

最后，附上GitHub地址:

https://github.com/ajbrock/BigGAN-PyTorch

還有論文：

Large Scale GAN Training for High Fidelity Natural Image Synthesis

https://arxiv.org/abs/1809.11096