谷歌宣布開源神經(jīng)網(wǎng)絡(luò)模型高效訓(xùn)練庫 GPipe

谷歌4日在博客中宣布開源大規(guī)模神經(jīng)網(wǎng)絡(luò)模型高效訓(xùn)練庫 GPipe，這是一款分布式機器學(xué)習(xí)庫，可以讓研究員在不調(diào)整超參數(shù)的情況下，部署更多的加速器以對大規(guī)模模型進行訓(xùn)練，有效擴展了模型性能。對此進行編譯如下。

深度神經(jīng)網(wǎng)絡(luò)（DNNs）推進諸多機器學(xué)習(xí)任務(wù)的進步，其中包括語音識別、視覺識別和語言處理等。BigGan、Bert 、GPT2.0 等最新成果表明，DNN 的模型越大，任務(wù)處理的表現(xiàn)就越好，而該結(jié)論也在過去的視覺識別任務(wù)中得到了驗證，表明模型大小與分類準確性之間存在很強的關(guān)聯(lián)性。舉個例子，2014 年 ImageNet 視覺識別挑戰(zhàn)賽的冠軍 GoogleNet 通過對 400 萬參數(shù)進行調(diào)整，最終取得 74.8 ％ 的精確度成績；僅僅過了三年，2017 年 ImageNet 挑戰(zhàn)賽冠軍 Squeeze-and-Excitation Networks 調(diào)整的參數(shù)便高達 1.458 億（36 倍以上），最終取得了 82.7％ 的精確度成績。與此對應(yīng)的是，市面的 GPU 內(nèi)存僅僅提高了 3 倍左右，目前最先進的圖像模型早已達到云 TPUv2 可用內(nèi)存的極限。因此，我們迫切需要一種高效、可擴展的基礎(chǔ)設(shè)施，以實現(xiàn)大規(guī)模的深度學(xué)習(xí)訓(xùn)練，并克服當前的加速器內(nèi)存受限問題。

最具有代表性的圖像分類模型體現(xiàn)出了 ImageNet 精確度結(jié)果與模型大小的強關(guān)聯(lián)性

在《GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism》一文中，我們展示了如何通過流水并行技術(shù)（pipeline parallelism）對 DNN 的訓(xùn)練進行擴展以克服這一限制。簡單來說，GPipe 是一款分布式機器學(xué)習(xí)庫，基于同步隨機梯度下降與流水并行技術(shù)進行模型訓(xùn)練，適用于由多個連續(xù)層組成的任意 DNN。最重要的是，GPipe 可以讓研究員在不調(diào)整超參數(shù)的情況下，部署更多的加速器以訓(xùn)練大規(guī)模模型，由此有效擴展了性能。為了證明 GPipe 的能力，我們在 Google Cloud TPUv2s 上訓(xùn)練了一個具有 5.57 億模型參數(shù)、 480 x 480 輸入圖像尺寸的 AmoebaNet-B。該模型在多個流行數(shù)據(jù)集上表現(xiàn)良好，取得的成就包括：single-crop ImageNet 的精確度提高至 84.3％、 CIFAR-10 的精確度提高至 99％、CIFAR-100 的精確度提高至 91.3％。

從小批次至微批次

目前存在兩種標準方法可以對中等規(guī)模的 DNN 模型進行加速。數(shù)據(jù)并行方法（The data parallelism）可以納入更多的機器，并將輸入的數(shù)據(jù)區(qū)分開來。另一種方法則是將模型置于加速器上（比如 GPU 或 TPU）——這些加速器的特殊硬件可加速模型的訓(xùn)練進程。然而加速器卻面臨著內(nèi)存與主機通信帶寬兩方面受限的問題。因此，通過將模型進行分區(qū)，并根據(jù)分區(qū)配置相應(yīng)的加速器，模型并行技術(shù)可以讓我們在加速器上訓(xùn)練更大規(guī)模的 DNN 模型。由于 DNN 存在順序性，這種策略最后可能變成計算期間只有一個加速器處于活躍狀態(tài)，未能將加速器的計算能力充分利用起來。此外，標準的數(shù)據(jù)并行技術(shù)只允許在多個加速器上同時訓(xùn)練具有不同輸入數(shù)據(jù)的相同模型，卻無法提升加速器所能支持的最大模型規(guī)模。

為了實現(xiàn)跨加速器的高效訓(xùn)練，GPipe 先按照加速器對模型進行劃分，然后自動將小批次的訓(xùn)練示例拆分為更小的微批次。通過在微批次中執(zhí)行流水管理，加速器得以并行運行。此外，梯度將在微批次中持續(xù)累積，以免分區(qū)的數(shù)量影響到模型的質(zhì)量。

上圖：由于網(wǎng)絡(luò)存在順序性，模型并行策略導(dǎo)致嚴重的利用不充分問題。 每次都只有一個加速器處于活動狀態(tài)。

下圖：GPipe 將輸入的小批次拆成更小的微批次，使不同的加速器可以同時在單獨的微批次上運作。

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

GPipe 會對模型參數(shù)的內(nèi)存分配進行最大化處理。我們在每個 TPUv2 均配有 8 個加速器核心以及 64 GB 內(nèi)存（每個加速器 8 GB）的云 TPUv2 上做了相關(guān)實驗。如果沒有 GPipe，由于內(nèi)存的限制，單個加速器頂多只能訓(xùn)練 8200 萬個模型參數(shù)。而通過反向傳播以及批量分割技術(shù)進行重新計算的 GPipe ，成功將中間激活內(nèi)存從 6.26 GB 減至 3.46GB，由此實現(xiàn)單個加速器上訓(xùn)練 3.18 億個參數(shù)的成果。此外，我們還發(fā)現(xiàn)在流水并行技術(shù)的作用下，模型的最大規(guī)模與分區(qū)數(shù)量成正比，正如事前所預(yù)料的那樣?？偟膩碚f，GPipe 使 AmoebaNet 能在云 TPUv2 的 8 個加速器上納入 18 億個參數(shù)，比起之前高出了 25 倍。

為了測試模型的效率，我們研究了 GPipe 對 AmoebaNet-D 模型吞吐量的影響情況。由于訓(xùn)練過程需要至少兩個加速器以適應(yīng)模型尺寸，因此我們只能對沒有實施流水并行技術(shù)的兩個分區(qū)案例的加速情況進行觀察。我們發(fā)現(xiàn)訓(xùn)練過程存在近乎線性的加速效果。與兩個分區(qū)案例相比，將模型分布在四倍數(shù)量的加速器上能有效實現(xiàn) 3.5 倍的加速效果。我們的實驗均使用了云 TPUv2，但我們了解到最新的的云 TPUv3 由于每個 TPUv3 均配備了 16 個加速器核心以及 256 GB（每個加速器 16 GB），因此擁有更理想的表現(xiàn)性能。當我們在所有 16 個加速器上對模型進行分發(fā)，GPipe 能讓基于 1024-token 句子的 80 億參數(shù) Transformer 語言模型的訓(xùn)練速度提高 11 倍。

通過 GPipe 對 AmoebaNet-D 進行加速。這種模型不適用于單個加速器。naive-2 基線是將模型拆分為兩個分區(qū)，最終取得的性能結(jié)果。Pipeline-k 對應(yīng)的是 GPipe 將模型分成帶有 k 個加速器的 k 個分區(qū)的最終性能結(jié)果。

在無需更改超參數(shù)的情況下，GPipe 還能通過使用更多加速器來擴展訓(xùn)練結(jié)果。因此，它可以與數(shù)據(jù)并行技術(shù)相結(jié)合，通過互補的方式使用更多的加速器來擴展神經(jīng)網(wǎng)絡(luò)的訓(xùn)練規(guī)模。

精準度測試

我們試圖通過 GPipe 證明，對現(xiàn)有的神經(jīng)網(wǎng)絡(luò)進行擴展，可以實現(xiàn)更理想的模型質(zhì)量。為此，我們在 ImageNet ILSVRC-2012 數(shù)據(jù)集上訓(xùn)練一個模型參數(shù)為 5.57 億、輸入圖像尺寸為 480 x 480 的 AmoebaNet-B。該網(wǎng)絡(luò)被分為 4 個分區(qū)，在模型與數(shù)據(jù)上執(zhí)行了并行訓(xùn)練程序。該巨型模型在沒有任何外部數(shù)據(jù)的情況下，最終達到了最先進的 84.3％ top-1 / 97％ top-5 的single-crop 驗證準確度結(jié)果。這說明大型的神經(jīng)網(wǎng)絡(luò)不僅適用于 ImageNet 等數(shù)據(jù)集，還能通過遷移學(xué)習(xí)的方式作用于其他數(shù)據(jù)集。事實證明，更好的 ImageNet 模型擁有更理想的傳輸效果。我們在 CIFAR10 和 CIFAR100 數(shù)據(jù)集上進行了遷移學(xué)習(xí)實驗。我們的巨型模型成功將 CIFAR-10 的精確度提高至到 99％、CIFAR-100 的精確度提高到 91.3％。

結(jié)論

當下許多機器學(xué)習(xí)應(yīng)用（如自動駕駛和醫(yī)學(xué)成像）得以持續(xù)發(fā)展并取得成功的原因，在于實現(xiàn)了盡可能高的模型精確度。然而這也意味著我們需要構(gòu)建一個更大、更復(fù)雜的模型，我們很高興能夠為研究社區(qū)提供 GPipe，我們希望未來它可以成為高效訓(xùn)練大規(guī)模 DNN 的基礎(chǔ)設(shè)施。