BigGan、Bert、GPT 2.0 等模型的出現(xiàn)似乎表明，DNN 模型越大，其準(zhǔn)確率越高。但芯片內(nèi)存的增長速度似乎無法滿足這種“暴力美學(xué)”的需要。為此，谷歌提出并開源了一種新的方法——GPipe，利用管道并行化擴展 DNN 訓(xùn)練以克服這種局限。它可以讓研究者輕松部署更多的加速器，以訓(xùn)練更大的模型，并在不調(diào)整超參數(shù)的前提下實現(xiàn)性能擴展。

深度神經(jīng)網(wǎng)絡(luò)（DNN）推動了許多機器學(xué)習(xí)任務(wù)的發(fā)展，包括語音識別、視覺識別、語言處理。BigGan、Bert、GPT 2.0取得的近期進(jìn)展表明，DNN模型越大，其在任務(wù)中的表現(xiàn)越好。視覺識別領(lǐng)域過去取得的進(jìn)展也表明，模型大小和分類準(zhǔn)確率之間存在很強的關(guān)聯(lián)。例如，2014年ImageNet視覺識別挑戰(zhàn)賽的冠軍GoogleNet以400萬的參數(shù)取得了74.8%的top-1準(zhǔn)確率，但僅僅過了三年，冠軍的寶座就被Squeeze-and-ExcitationNetworks搶去，后者以1.458億（前者的36倍還多）的參數(shù)量取得了82.7%的top-1準(zhǔn)確率。然而，在這段時間里，GPU的內(nèi)存只提高了3倍左右，當(dāng)前最優(yōu)的圖像模型卻已經(jīng)達(dá)到了谷歌云 TPUv2的可用內(nèi)存。因此，我們急需一個能夠?qū)崿F(xiàn)大規(guī)模深度學(xué)習(xí)并克服當(dāng)前加速器內(nèi)存局限的可擴展高效架構(gòu)。

近期有代表性的圖像分類模型的ImageNet準(zhǔn)確率和模型大小之間的強關(guān)聯(lián)。

谷歌在論文《GPipe: Efficient Training of GiantNeural Networks using Pipeline Parallelism》中，展示了利用管道并行化（pipeline parallelism）擴展DNN訓(xùn)練以克服這種局限。GPipe是一個分布式機器學(xué)習(xí)庫，它使用同步隨機梯度下降和管道并行化進(jìn)行訓(xùn)練，可以應(yīng)用到包含多個序列層的任意DNN中。重要的是，GPipe可以讓研究者輕松部署更多的加速器，以訓(xùn)練更大的模型，并在不調(diào)整超參數(shù)的前提下實現(xiàn)性能擴展。為了展示GPipe的有效性，谷歌在谷歌云TPUv2上訓(xùn)練了AmoebaNet-B模型，該模型具備5.57億個參數(shù)，輸入圖像大小為480 x 480。這一模型在多個流行的數(shù)據(jù)集上表現(xiàn)良好，比如將ImageNet上的single-crop準(zhǔn)確率提高到84.3%，將CIFAR-10上的準(zhǔn)確率提高到99%，CIFAR-100準(zhǔn)確率提高到91.3%。該核心GPipe庫已經(jīng)在Lingvo框架下開源。

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

從小批量到微小批量

加速中型DNN模型有兩種標(biāo)準(zhǔn)方法。第一種是數(shù)據(jù)并行化方法，它利用更多機器將輸入數(shù)據(jù)分散到這些機器上。另一種方法是將模型移到加速器上，如GPU 或 TPU，這些加速器具備特殊的硬件可以加速模型訓(xùn)練。但是，加速器內(nèi)存有限、通信帶寬也有限。因此，要在加速器上訓(xùn)練大型DNN模型需要模型并行化，將模型分割成多個部分，然后將不同的部分分配到不同的加速器。但是由于DNN的序列本質(zhì)，這一策略可能導(dǎo)致計算期間只有一個加速器處于活躍狀態(tài)，而這會導(dǎo)致無法充分利用加速器的計算能力。另外，標(biāo)準(zhǔn)數(shù)據(jù)并行化方法允許同一個模型在多個加速器上對不同的輸入數(shù)據(jù)執(zhí)行并行訓(xùn)練，但是這無法增加每個加速器可以支持的最大模型大小。

為了在多個加速器上也能進(jìn)行高效的模型訓(xùn)練，GPipe將模型分割并分配給不同的加速器，將小批量訓(xùn)練樣本自動分割成更小的批量（微小批量）。通過在微小批量樣本上管道化整個執(zhí)行過程，加速器可以實現(xiàn)并行運行。此外，各微小批量上的梯度可以一直累加，這樣分區(qū)數(shù)量就不會影響到模型質(zhì)量。

上：由于DNN網(wǎng)絡(luò)的序列本質(zhì)，初始的模型并行化策略導(dǎo)致計算能力無法充分利用，每次只有一個加速器處于活躍狀態(tài)。下：GPipe 將輸入小批量分割成更小的批量，使得不同的加速器可以同時處理各自分配到的微小批量樣本。

最大化內(nèi)存和效率

GPipe 最大化了模型參數(shù)的內(nèi)存分配。谷歌在云 TPUv2上進(jìn)行試驗，每個芯片上有8個加速器內(nèi)核和64GB的內(nèi)存（每個加速器8GB內(nèi)存）。沒有TPUv2的情況下，由于內(nèi)存限制，單個加速器最多訓(xùn)練8200萬模型參數(shù)。由于反向傳播和批量分割中的重復(fù)計算，GPipe 將中間激活內(nèi)存從6.26GB降至3.46GB，使得單個加速器上可以訓(xùn)練3.18億個參數(shù)。谷歌還發(fā)現(xiàn)，如預(yù)期所料，在管道并行化的情況下，最大模型大小與分區(qū)數(shù)量成正比。有了GPipe，AmoebaNet 能夠在云TPUv2的8個加速器上加入18億參數(shù)，是沒有GPipe時的25倍。

為了測試效率，谷歌衡量了GPipe對AmoebaNet-D模型吞吐量的影響。因為訓(xùn)練至少需要兩個加速器來適應(yīng)模型大小，谷歌衡量了在兩個分區(qū)但沒有管道并行化的naive情況下的加速，發(fā)現(xiàn)訓(xùn)練過程中幾乎是線性加速。與具有兩個分區(qū)的naive方法相比，將模型分布在四倍的加速器上實現(xiàn)了3.5倍的加速。雖然本文所有的實驗都使用了云TPUv2，但谷歌發(fā)現(xiàn)當(dāng)前可用的云TPUv3會有更好的性能，每個芯片有16個加速器內(nèi)核和256GB內(nèi)存（每個加速器16GB）。當(dāng)在所有16個加速器上分布模型時，GPipe使得具有80億參數(shù)的Transformer 語言模型在1024-token語句上的加速達(dá)到11倍。

使用GPipe加速AmoebaNet-D。該模型不適合加速器?；€naive-2是模型被分割成兩個分區(qū)時native分區(qū)方法的性能。Pipeline-k是使用K個加速器將模型分割成K個分區(qū)時GPipe的性能。

GPipe還可以通過使用更多的加速器而不改變參數(shù)來擴大訓(xùn)練規(guī)模。因此，它可以與數(shù)據(jù)并行相結(jié)合，以互補的方式使用更多的加速器來擴大神經(jīng)網(wǎng)絡(luò)訓(xùn)練規(guī)模。

測試準(zhǔn)確率

谷歌使用GPipe 來驗證一個假設(shè)，即擴大現(xiàn)有神經(jīng)網(wǎng)絡(luò)能夠提高模型質(zhì)量。谷歌在 ImageNet ILSVRC-2012 數(shù)據(jù)集上訓(xùn)練了一個AmoebaNet-B模型，其模型參數(shù)為5.57億，輸入圖像大小為480 x 480。該網(wǎng)絡(luò)被分為四個分區(qū)，且在模型和數(shù)據(jù)上應(yīng)用了并行訓(xùn)練過程。這個巨大的模型在沒有任何外部數(shù)據(jù)的情況下達(dá)到了當(dāng)前最先進(jìn)的84.3% top-1 / 97% top-5 single-crop驗證準(zhǔn)確率。大型神經(jīng)網(wǎng)絡(luò)不僅適用于如ImageNet這樣的數(shù)據(jù)集，并且通過遷移學(xué)習(xí)還可適用于其他數(shù)據(jù)集。已有研究證明，更好的ImageNet模型遷移效果更好。谷歌在CIFAR10 和CIFAR100數(shù)據(jù)集上進(jìn)行遷移學(xué)習(xí)實驗。其巨大的模型將 CIFAR-10的準(zhǔn)確率提高到99%，將CIFAR-100的準(zhǔn)確率提高到91.3%。

表2：單個AmoebaNet-B(6, 512) 模型在ImageNet ILSVRC 2012驗證數(shù)據(jù)集上的single-crop分類準(zhǔn)確率與其他已發(fā)布最優(yōu)模型的對比結(jié)果。

表3：AmoebaNet-B (6,512) 的遷移學(xué)習(xí)結(jié)果，輸入圖像大小為480×480，測試時使用的是single crop。這些結(jié)果均為5次精調(diào)運行后的平均值。

總結(jié)

很多機器學(xué)習(xí)實際應(yīng)用的持續(xù)發(fā)展和成功依賴于盡可能高的準(zhǔn)確率，如自動駕駛和醫(yī)療成像。而這通常需要構(gòu)建更大型、更復(fù)雜的模型，谷歌很高興將GPipe帶給更廣大的社區(qū)，并希望它能夠有效訓(xùn)練大型深度神經(jīng)網(wǎng)絡(luò)。

論文：GPipe: Efficient Training of GiantNeural Networks using Pipeline Parallelism

論文鏈接：https://arxiv.org/pdf/1811.06965.pdf

摘要：GPipe是一個可擴展管道并行化庫，可以訓(xùn)練超大型的深度神經(jīng)網(wǎng)絡(luò)。它將網(wǎng)絡(luò)的層分割成多個部分并分配給不同的加速器，將執(zhí)行過程管道化以達(dá)到最高的硬件利用率。例如，在8個加速器上應(yīng)用分割技術(shù)，則GPipe支持25倍大小的神經(jīng)網(wǎng)絡(luò)的訓(xùn)練，這充分展示了其可擴展性。它還保證了無論分區(qū)數(shù)量多少，梯度仍是一致的。在對模型參數(shù)幾乎不做更改的情況下，它獲得了幾乎是線性的加速：加速器數(shù)量是原來的4倍時，同一模型的訓(xùn)練速度是原來的3.5倍。我們訓(xùn)練了一個具備5.57億參數(shù)的AmoebaNet模型，在 ImageNet 2012 數(shù)據(jù)集上達(dá)到了新的SOTA結(jié)果：84.3%的 top-1 準(zhǔn)確率和97.0% 的 top-5 準(zhǔn)確率。最后，我們將這個學(xué)到的模型應(yīng)用到多個流行圖像分類數(shù)據(jù)集上，同樣獲得了有競爭力的結(jié)果，比如在CIFAR-10數(shù)據(jù)集上達(dá)到了99%的準(zhǔn)確率，在CIFAR-100數(shù)據(jù)集上達(dá)到了91.3%的準(zhǔn)確率。