模型參數太多怎麼辦?用Google高效訓練庫GPipe啊

尋夢新聞LINE@每日推播熱門推薦文章,趣聞不漏接❤️

加入LINE好友

曉查 發自 凹非寺量子位 報導 | 公眾號 QbitAI

BigGan、BERT和GPT-2的出現表明,越大的神經網路模型可以帶來越好的性能。其中能寫出各種假新聞的GPT-2參數就有15億個之多。

如果你也在訓練一個參數超多的模型,那麼現在有個好消息要告訴你。

Google在最近公布的Lingvo框架下開源了GPipe。它讓研究人員不用調整超參數,就能用更多加速核心來訓練更大的模型和提升性能。

GPipe適用於多個連續層組成的任何DNN,而且訓練模型的大小與分區數成正比。

參數越多,效果越好

在計算機視覺圖像分類任務上,模型大小和分類準確性之間存在很強的相關性。

模型參數太多怎麼辦?用Google高效訓練庫GPipe啊 家居 第1張

不同圖像分類模型在ImageNet上的表現,顯示出準確度和參數數量的強相關性

從2014年到2017年,ImageNet視覺識別的冠軍的準確度從74.8%提升到了82.7%,模型參數增加了36倍。

而在這三年間,GPU顯存僅僅增加了3倍,因此迫切需要一種高效的神經網路訓練庫來解決參數太多的問題。

核心越多,參數越多

GPipe就可以減少參數模型對內存的占用。

Google研究人員在8核心64GB內存的雲TPUv2進行測試。如果沒有GPipe,由於內存的限制,每個加速核心只能訓練8200萬個參數。

使用GPipe後,活動內存從6.26 GB降低到3.46GB,可以在單個加速核心上訓練3.18億個參數,接近原來的4倍。

通過GPipe,AmoebaNet能夠在雲TPUv2上加入18億個參數,比沒有GPipe的情況下多25倍。

模型參數太多怎麼辦?用Google高效訓練庫GPipe啊 家居 第2張

Google將AmoebaNet-D模型分配到四倍數量的加速核心上,做到了3.5倍的加速

GPipe不僅僅速度更快,在訓練模型的質量上也有提高。

為了驗證這一點,Google訓練了另一種圖像分類算法AmoebaNet-B,它總共有5.57億個模型參數。

該算法在一些常用的數據集上表現很好,將single-crop ImageNet精度提高到了84.3%,在CIFAR-10上的精度為99%,在CIFAR-100上的精度為91.3%。

GPipe為何能做到高效

原則上我們可以通過並行計算在GPU或者TPU上訓練更大的DNN模型。但是由於DNN的順序性,這種方法可能導致在計算期間只有一個加速器處於活動狀態,不能充分利用設備的計算能力。

Google在GPipe使用了兩種AI訓練技術:一種是同步隨機梯度下降,另一種是並行pipeline,將上一步的輸出作為流傳輸到下一步。

Gpipe將模型劃分到不同的加速核心,並自動將mini-batch的訓練樣本分成更小的micro-batch,並在pipeline中運行,使TPU核心能夠並行操作。

此外,梯度一直在micro-batch中累積,因此分區數量不會影響模型質量。

模型參數太多怎麼辦?用Google高效訓練庫GPipe啊 家居 第3張

以上都是在雲TPUv2上的結果,實際上GPipe在雲TPUv3上有更好的性能,每個TPUv3都有16個加速器核心和256 GB(每個加速器16 GB)。用GPipe訓練80億參數的Transformer模型時,能做到11倍的加速。

開源地址:https://github.com/tensorflow/lingvo/blob/master/lingvo/core/gpipe.py