你也可以訓練超大神經網路！谷歌開源GPipe庫

BigGan、Bert、GPT 2.0 等模型的出現似乎表明，DNN 模型越大，其準確率越高。但晶片記憶體的增長速度似乎無法滿足這種“暴力美學”的需要。為此，谷歌提出並開源了一種新的方法——GPipe，利用管道並行化擴充套件 DNN 訓練以克服這種侷限。它可以讓研究者輕鬆部署更多的加速器，以訓練更大的模型，並在不調整超引數的前提下實現效能擴充套件。

深度神經網路（DNN）推動了許多機器學習任務的發展，包括語音識別、視覺識別、語言處理。BigGan、Bert、GPT 2.0取得的近期進展表明，DNN模型越大，其在任務中的表現越好。視覺識別領域過去取得的進展也表明，模型大小和分類準確率之間存在很強的關聯。例如，2014年ImageNet視覺識別挑戰賽的冠軍GoogleNet以400萬的引數取得了74.8%的top-1準確率，但僅僅過了三年，冠軍的寶座就被Squeeze-and-ExcitationNetworks搶去，後者以1.458億（前者的36倍還多）的引數量取得了82.7%的top-1準確率。然而，在這段時間裡，GPU的記憶體只提高了3倍左右，當前最優的影像模型卻已經達到了谷歌雲 TPUv2的可用記憶體。因此，我們急需一個能夠實現大規模深度學習並克服當前加速器記憶體侷限的可擴充套件高效架構。 

近期有代表性的影像分類模型的ImageNet準確率和模型大小之間的強關聯。

谷歌在論文《GPipe: Efficient Training of GiantNeural Networks using Pipeline Parallelism》中，展示了利用管道並行化（pipeline parallelism）擴充套件DNN訓練以克服這種侷限。GPipe是一個分散式機器學習庫，它使用同步隨機梯度下降和管道並行化進行訓練，可以應用到包含多個序列層的任意DNN中。重要的是，GPipe可以讓研究者輕鬆部署更多的加速器，以訓練更大的模型，並在不調整超引數的前提下實現效能擴充套件。為了展示GPipe的有效性，谷歌在谷歌雲TPUv2上訓練了AmoebaNet-B模型，該模型具備5.57億個引數，輸入影像大小為480 x 480。這一模型在多個流行的資料集上表現良好，比如將ImageNet上的single-crop準確率提高到84.3%，將CIFAR-10上的準確率提高到99%，CIFAR-100準確率提高到91.3%。該核心GPipe庫已經在Lingvo框架下開源。

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

從小批量到微小批量

加速中型DNN模型有兩種標準方法。第一種是資料並行化方法，它利用更多機器將輸入資料分散到這些機器上。另一種方法是將模型移到加速器上，如GPU 或 TPU，這些加速器具備特殊的硬體可以加速模型訓練。但是，加速器記憶體有限、通訊頻寬也有限。因此，要在加速器上訓練大型DNN模型需要模型並行化，將模型分割成多個部分，然後將不同的部分分配到不同的加速器。但是由於DNN的序列本質，這一策略可能導致計算期間只有一個加速器處於活躍狀態，而這會導致無法充分利用加速器的計算能力。另外，標準資料並行化方法允許同一個模型在多個加速器上對不同的輸入資料執行並行訓練，但是這無法增加每個加速器可以支援的最大模型大小。

為了在多個加速器上也能進行高效的模型訓練，GPipe將模型分割並分配給不同的加速器，將小批量訓練樣本自動分割成更小的批量（微小批量）。通過在微小批量樣本上管道化整個執行過程，加速器可以實現並行執行。此外，各微小批量上的梯度可以一直累加，這樣分割槽數量就不會影響到模型質量。

上：由於DNN網路的序列本質，初始的模型並行化策略導致計算能力無法充分利用，每次只有一個加速器處於活躍狀態。下：GPipe 將輸入小批量分割成更小的批量，使得不同的加速器可以同時處理各自分配到的微小批量樣本。

最大化記憶體和效率

GPipe 最大化了模型引數的記憶體分配。谷歌在雲 TPUv2上進行試驗，每個晶片上有8個加速器核心和64GB的記憶體（每個加速器8GB記憶體）。沒有TPUv2的情況下，由於記憶體限制，單個加速器最多訓練8200萬模型引數。由於反向傳播和批量分割中的重複計算，GPipe 將中間啟用記憶體從6.26GB降至3.46GB，使得單個加速器上可以訓練3.18億個引數。谷歌還發現，如預期所料，在管道並行化的情況下，最大模型大小與分割槽數量成正比。有了GPipe，AmoebaNet 能夠在雲TPUv2的8個加速器上加入18億引數，是沒有GPipe時的25倍。

為了測試效率，谷歌衡量了GPipe對AmoebaNet-D模型吞吐量的影響。因為訓練至少需要兩個加速器來適應模型大小，谷歌衡量了在兩個分割槽但沒有管道並行化的naive情況下的加速，發現訓練過程中幾乎是線性加速。與具有兩個分割槽的naive方法相比，將模型分佈在四倍的加速器上實現了3.5倍的加速。雖然本文所有的實驗都使用了雲TPUv2，但谷歌發現當前可用的雲TPUv3會有更好的效能，每個晶片有16個加速器核心和256GB記憶體（每個加速器16GB）。當在所有16個加速器上分佈模型時，GPipe使得具有80億引數的Transformer 語言模型在1024-token語句上的加速達到11倍。 

使用GPipe加速AmoebaNet-D。該模型不適合加速器。基線naive-2是模型被分割成兩個分割槽時native分割槽方法的效能。Pipeline-k是使用K個加速器將模型分割成K個分割槽時GPipe的效能。

GPipe還可以通過使用更多的加速器而不改變引數來擴大訓練規模。因此，它可以與資料並行相結合，以互補的方式使用更多的加速器來擴大神經網路訓練規模。

測試準確率

谷歌使用GPipe 來驗證一個假設，即擴大現有神經網路能夠提高模型質量。谷歌在 ImageNet ILSVRC-2012 資料集上訓練了一個AmoebaNet-B模型，其模型引數為5.57億，輸入影像大小為480 x 480。該網路被分為四個分割槽，且在模型和資料上應用了並行訓練過程。這個巨大的模型在沒有任何外部資料的情況下達到了當前最先進的84.3% top-1 / 97% top-5 single-crop驗證準確率。大型神經網路不僅適用於如ImageNet這樣的資料集，並且通過遷移學習還可適用於其他資料集。已有研究證明，更好的ImageNet模型遷移效果更好。谷歌在CIFAR10 和CIFAR100資料集上進行遷移學習實驗。其巨大的模型將 CIFAR-10的準確率提高到99%，將CIFAR-100的準確率提高到91.3%。

表2：單個AmoebaNet-B(6, 512) 模型在ImageNet ILSVRC 2012驗證資料集上的single-crop分類準確率與其他已釋出最優模型的對比結果。

表3：AmoebaNet-B (6,512) 的遷移學習結果，輸入影像大小為480×480，測試時使用的是single crop。這些結果均為5次精調執行後的平均值。 

總結

很多機器學習實際應用的持續發展和成功依賴於儘可能高的準確率，如自動駕駛和醫療成像。而這通常需要構建更大型、更復雜的模型，谷歌很高興將GPipe帶給更廣大的社群，並希望它能夠有效訓練大型深度神經網路。

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

論文連結：https://arxiv.org/pdf/1811.06965.pdf

摘要：GPipe是一個可擴充套件管道並行化庫，可以訓練超大型的深度神經網路。它將網路的層分割成多個部分並分配給不同的加速器，將執行過程管道化以達到最高的硬體利用率。例如，在8個加速器上應用分割技術，則GPipe支援25倍大小的神經網路的訓練，這充分展示了其可擴充套件性。它還保證了無論分割槽數量多少，梯度仍是一致的。在對模型引數幾乎不做更改的情況下，它獲得了幾乎是線性的加速：加速器數量是原來的4倍時，同一模型的訓練速度是原來的3.5倍。我們訓練了一個具備5.57億引數的AmoebaNet模型，在 ImageNet 2012 資料集上達到了新的SOTA結果：84.3%的 top-1 準確率和97.0% 的 top-5 準確率。最後，我們將這個學到的模型應用到多個流行影像分類資料集上，同樣獲得了有競爭力的結果，比如在CIFAR-10資料集上達到了99%的準確率，在CIFAR-100資料集上達到了91.3%的準確率。

原文連結：https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html