Github上放出了只需4-8塊GPU就能訓(xùn)練的“改進版”BigGAN模型代碼

機器學(xué)習(xí)模型訓(xùn)練成本往往令普通人倍感頭疼，動輒幾十上百塊泰坦，別說買，就是租都肉疼。近日，BigGAN作者之一在Github上放出了只需4-8塊GPU就能訓(xùn)練的“改進版”BigGAN模型代碼，可以說是窮人的福音。新模型使用PyTorch實現(xiàn)。

機器學(xué)習(xí)模型訓(xùn)練是一個耗時費力的過程，而且隨著人們對模型性能要求的提升，訓(xùn)練模型需要的計算力正以驚人的速度增長，堆疊高性能GPU進行數(shù)據(jù)訓(xùn)練幾乎是唯一選擇，動輒幾十塊上百塊的泰坦，搞的地主家也沒有余糧。

BigGAN效果拔群，但訓(xùn)練成本同樣讓人望而卻步，想自己搞？先摸摸錢包再說。

現(xiàn)在，BigGAN原作者之一Andrew Brock在Github上放出了只需4-8塊GPU就能訓(xùn)練的新版BigGAN，想窮人之所想，急窮人之所急，可以說是非常親民了。新模型使用的是PyTorch，而不是TF。

下面一起看看這個新模型的具體介紹，以下內(nèi)容來自Github上的簡介。

本資源包含由Andrew Brock，JeffDonahue和Karen Simonyan進行的大規(guī)模GAN高保真自然圖像合成訓(xùn)練的BigGAN，只需4-8塊 GPU的訓(xùn)練代碼。

本段代碼由Andy Brock和Alex Andonian編寫。

運行環(huán)境和條件

PyTorch 1.0.1

tqdm，numpy，scipy和h5py

ImageNet訓(xùn)練集

首先，可以選擇準(zhǔn)備目標(biāo)數(shù)據(jù)集的預(yù)處理HDF5版本，以實現(xiàn)更快的輸入輸出。之后需要計算FID所需的Inception時刻。這些都可以通過修改和運行以下代碼來完成

shscripts / utils / prepare_data.sh

默認(rèn)情況下，ImageNet訓(xùn)練集被下載到此目錄中的根文件夾中，并將以128x128像素分辨率準(zhǔn)備緩存的HDF5。

在scripts文件夾中，有多個bash腳本可以訓(xùn)練具有不同批量大小的BigGAN。假設(shè)您無法訪問完整的TPU pod，因此通過梯度累積（在多個小批量下進行梯度平均，并且僅在N次累積后執(zhí)行優(yōu)化程序步驟），以此形式表示大批量。

默認(rèn)情況下，可以使用launch_BigGAN_bs256x8.sh腳本訓(xùn)練一個全尺寸的BigGAN模型，批大小為256和8個梯度累積，總批量為2048。在8張V100上進行全精度訓(xùn)練（無張量），訓(xùn)練需要15天，期間共進行約150k次迭代。

首先需要確定設(shè)置可以支持的最大批量大小。這里提供的預(yù)訓(xùn)練模型是在8個V100上（每個顯存16GB ）上訓(xùn)練的，這個配置可以支持比默認(rèn)使用的B1S256稍多一些的載荷。一旦確定了這一點，就應(yīng)該修改腳本，使批大小乘以梯度累積的數(shù)量等于所需的總批量大?。˙igGAN默認(rèn)為2048）。

另外，此腳本使用--load_in_memarg，將整個（最大支持64GB）的I128.hdf5文件加載到RAM中，以加快數(shù)據(jù)的加載速度。如果沒有足夠的RAM做硬件支持（可能需要96GB以上的RAM），請刪除此參數(shù)。

度量標(biāo)準(zhǔn)和抽樣

在訓(xùn)練期間，腳本將輸出帶有訓(xùn)練指標(biāo)和測試指標(biāo)的日志，同時保存模型權(quán)重和優(yōu)化程序參數(shù)的多個副本（前者保存最近的2個，后者保存5個最高得分），并且每次保存權(quán)重時將生成樣本和插值。 logs文件夾包含處理這些日志的腳本，并使用MATLAB繪制結(jié)果。

訓(xùn)練之后，可以使用sample.py生成其他樣本和插值，使用不同的截斷值，批量大小，站立統(tǒng)計累積次數(shù)等進行測試。有關(guān)示例，請參閱sample_BigGAN_bs256x8.sh腳本。

默認(rèn)情況下，所有內(nèi)容都保存在weights/samples/logs/data文件夾中，這些文件夾設(shè)置與此repo位于同一文件夾中?？梢允褂?-base_root參數(shù)將所有這些指向不同的基本文件夾，或者使用各自的參數(shù)（例如--logs_root）選擇每個基礎(chǔ)文件夾的特定位置。

此代碼中包含了運行BigGAN-deep的腳本，但還沒有完全訓(xùn)練使用它們的模型，因此用戶可以視作這些模型尚未測試過。此外，我代碼中還包括在CIFAR上運行模型的腳本，以及在ImageNet上運行SA-GAN（包括EMA）和SN-GAN的腳本。

SA-GAN代碼假設(shè)用戶配置在4張TitanX（或等同于該配置的GPU RAM），并且將以批量大小為128以及2個梯度累積運行。

關(guān)于初始度量標(biāo)準(zhǔn)的重要說明

本資源使用PyTorch內(nèi)置的初始網(wǎng)絡(luò)來計算IS和FID分?jǐn)?shù)。這些分?jǐn)?shù)與使用Tensorflow官方初始代碼獲得的分?jǐn)?shù)不同，僅用于監(jiān)控目的。使用--sample_npz參數(shù)在模型上運行sample.py，然后運行inception_tf13來計算實際的TensorFlow IS。請注意，需要安裝TensorFlow 1.3或更早版本，因為1.4或更高版本會破壞原始的IS代碼。

預(yù)訓(xùn)練模型

PyTorch初始分?jǐn)?shù)和FID分?jǐn)?shù)

我們引入了兩個預(yù)訓(xùn)練模型檢查點（使用G，D，G的EMA副本，優(yōu)化器和狀態(tài)dict）：

主要檢查點是在 128x128 ImageNet 圖像上訓(xùn)練的 BigGAN，該模型使用 BS256 和 8 梯度累積，并在崩潰前實現(xiàn)，其 TF Inception Score 為 97.35 +/- 1.79，

詳見：

https://drive.google.com/open?id=1nAle7FCVFZdix2—ks0r5JBkFnKw8ctW

第一個模型（100k G iters）的早期檢查點，性能更高，在崩潰之前實現(xiàn)，可能更容易微調(diào)。

詳見：

https://drive.google.com/open?id=1dmZrcVJUAWkPBGza_XgswSuT-UODXZcO。

另外，使用Places-365數(shù)據(jù)集的預(yù)訓(xùn)練模型即將推出。

此repo還包含用于將原始TF HubBigGAN 生成器權(quán)重的PyTorch的移植腳本。有關(guān)更多詳細(xì)信息，請參閱TFHub文件夾中的腳本。

使用自己的數(shù)據(jù)集或創(chuàng)建新的訓(xùn)練函數(shù)微調(diào)模型

如果想恢復(fù)中斷訓(xùn)練或微調(diào)預(yù)訓(xùn)練模型，請在運行相同的啟動腳本，添加--resume參數(shù)。實驗名稱是由訓(xùn)練配置自動生成的，但如果希望使用修改后的優(yōu)化器設(shè)置微調(diào)模型，可以使用--experiment_namearg進行文件名的覆蓋。

要準(zhǔn)備自己的數(shù)據(jù)集，需要將其添加到datasets.py并修改utils.py中的convenience dicts，以獲得數(shù)據(jù)集的相應(yīng)元數(shù)據(jù)。在prepare_data.sh中重復(fù)此過程（也可以選擇生成HDF5預(yù)處理副本，并計算FID的Inception Moments）。

默認(rèn)情況下，訓(xùn)練腳本將保存初始分?jǐn)?shù)最高的前5個檢查點。對于ImageNet以外的數(shù)據(jù)集，初始分?jǐn)?shù)可能是一種非常差的質(zhì)量標(biāo)準(zhǔn)，可以使用--which_bestFID來代替。

要使用自己的訓(xùn)練函數(shù)（如訓(xùn)練BigVAE）：修改train_fns.GAN_training_function或在if config['which_train_fn'] =='GAN'之后添加新的訓(xùn)練函數(shù)。

本模型的主要亮點

本資源庫提供完整的訓(xùn)練和指標(biāo)日志以供參考。重現(xiàn)論文過程中最困難的事情之一就是檢查訓(xùn)練早期的記錄日志是否規(guī)整，特別是在訓(xùn)練時間長達數(shù)周的情況下。希望這將有助于未來的工作。

本資源庫包括一個加速的FID計算 - 原始的scipy版本可能需要超過10分鐘來計算矩陣sqrt，此版本使用加速的PyTorch版本，計算時間不到1秒。

本資源用了一種加速、低內(nèi)存消耗的正交寄存器實現(xiàn)。默認(rèn)情況下，只計算最大奇異值（譜范數(shù)），但本段代碼通過 —num_G_SVs 參數(shù)支持了更多 SV 的計算。

本模型與原始BigGAN之間的主要區(qū)別

我們使用來自SA-GAN的優(yōu)化器設(shè)置（G_lr= 1e-4，D_lr = 4e-4，num_D_steps= 1，與BigGAN的設(shè)置不同（G_lr = 5e-5，D_lr = 2e-5，num_D_steps = 2）。雖然這樣犧牲了些許性能，但這是削減訓(xùn)練時間的第一步。

默認(rèn)情況下，本資源不使用Cross-Replica BatchNorm（又名Synced BatchNorm）。本資源嘗試的兩種變體與內(nèi)置的BatchNorm具有略微不同的梯度（盡管是相同的前向傳遞），可以滿足訓(xùn)練要求。

梯度累積意味著需要更頻繁地更新SV估計值和BN統(tǒng)計量（頻度增加了8倍）。這意味著BN統(tǒng)計數(shù)據(jù)更接近于常設(shè)統(tǒng)計數(shù)據(jù)，而且奇異值估計往往更準(zhǔn)確。因此，在測試模式下默認(rèn)使用G來衡量指標(biāo)（使用BatchNorm運行統(tǒng)計估算，而不是像文件中那樣計算常設(shè)統(tǒng)計數(shù)據(jù)）。

我們?nèi)匀恢С殖ＴO(shè)統(tǒng)計信息（具體見sample.sh腳本）。這也可能導(dǎo)致早期累積的梯度變得過時，但在實踐中這已經(jīng)不再是個問題。

目前給出的預(yù)訓(xùn)練模型未經(jīng)過正交正則化訓(xùn)練。似乎增加了模型由于截斷變得不可修復(fù)的可能性，但本資源庫中給出特定模型似乎格外好運，沒有碰到這種情況。不過，我們還是提供兩個經(jīng)過高度優(yōu)化（快速和最小內(nèi)存消耗）的正交寄存器實現(xiàn)，直接計算正交寄存器梯度。

Github資源地址：

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