隨著技術(shù)、算力的發(fā)展，在 ImageNet 上訓(xùn)練 ResNet-50 的速度被不斷刷新。2018 年 7 月，騰訊機(jī)智機(jī)器學(xué)習(xí)平臺(tái)團(tuán)隊(duì)在 ImageNet 數(shù)據(jù)集上僅用 6.6 分鐘就訓(xùn)練好 ResNet-50，創(chuàng)造了 AI 訓(xùn)練世界紀(jì)錄；一周前，壕無(wú)人性的索尼用 2176 塊 V100 GPU 將這一紀(jì)錄縮短到了 224 秒；如今，這一紀(jì)錄再次被谷歌刷新……

深度學(xué)習(xí)非常依賴(lài)于硬件條件，它是一個(gè)計(jì)算密集型的任務(wù)。硬件供應(yīng)商通過(guò)在大型計(jì)算集群中部署更快的加速器來(lái)做出更快的相應(yīng)。在 petaFLOPS（運(yùn)算能力單位，每秒千萬(wàn)億次浮點(diǎn)數(shù)運(yùn)算）規(guī)模的設(shè)備上訓(xùn)練深度學(xué)習(xí)模型需要同時(shí)面臨算法和系統(tǒng)軟件兩方面的挑戰(zhàn)。Google 于近日推出了一種大規(guī)模計(jì)算集群的圖像分類(lèi)人物訓(xùn)練解決方案，相關(guān)論文發(fā)表于 Arxiv：Image Classification at Supercomputer Scale。本文的作者使用 Google TPU v3 Pod 訓(xùn)練 ResNet-50，在識(shí)別率沒(méi)有降低的情況下，僅使用了 2.2 分鐘。

背景

深度神經(jīng)網(wǎng)絡(luò)的成功應(yīng)用與發(fā)展離不開(kāi)瘋狂增長(zhǎng)的算力，在許多領(lǐng)域，深度學(xué)習(xí)的發(fā)展可以說(shuō)是由硬件驅(qū)動(dòng)的。在深度網(wǎng)絡(luò)的訓(xùn)練過(guò)程中，最關(guān)鍵的部分就是使用隨機(jī)梯度下降算法（SGD）優(yōu)化網(wǎng)絡(luò)權(quán)重。通常情況下，模型需要使用 SGD 在一個(gè)數(shù)據(jù)集上進(jìn)行多次的便利才能達(dá)到收斂。在整個(gè)過(guò)程中，浮點(diǎn)數(shù)運(yùn)算能力顯得至關(guān)重要。例如，在 ImageNet 數(shù)據(jù)庫(kù)上訓(xùn)練 ResNet-50 模型，遍歷一次數(shù)據(jù)庫(kù)需要 3.2 萬(wàn)萬(wàn)億次浮點(diǎn)數(shù)運(yùn)算。而使模型達(dá)到收斂，通常需要遍歷 90 次數(shù)據(jù)庫(kù)。

盡管硬件加速設(shè)備（例如 GPU、TPU）已經(jīng)加快了迭代的次數(shù)，使用單個(gè)加速設(shè)備在大規(guī)模數(shù)據(jù)庫(kù)訓(xùn)練大型的神經(jīng)網(wǎng)絡(luò)仍然需要幾個(gè)小時(shí)或數(shù)天的時(shí)間。最常見(jiàn)的加速方法便是通過(guò)分布式的 SGD 算法使用多個(gè)設(shè)備并行訓(xùn)練，將每個(gè) mini-batch 分布在多個(gè)相同的加速設(shè)備上。

以往大家都喜歡用異步分布式 SGD 算法在將多個(gè)線程聯(lián)合起來(lái)進(jìn)行訓(xùn)練，但是近期的一些工作發(fā)現(xiàn)，異步分布式 SGD 算法優(yōu)化的模型在收斂程度和驗(yàn)證準(zhǔn)確率方面都不如同步分布式 SGD 訓(xùn)練出的模型。但是，為了保證在提速的同時(shí)模型的質(zhì)量不會(huì)有所損失，在使用同步分布式 SGD 算法的過(guò)程中，會(huì)遇到很多技術(shù)和硬件方面的瓶頸，作者總結(jié)出以下幾點(diǎn)：

模型的準(zhǔn)確率依賴(lài)于全局的 batch size 和計(jì)算集群中每個(gè)節(jié)點(diǎn)的 batch size。

在加速設(shè)備計(jì)算能力足夠高時(shí)，CPU 向 GPU 等專(zhuān)用設(shè)備的輸入過(guò)程成為了訓(xùn)練過(guò)程中的瓶頸。

使用同步分布式 SGD 算法需要大規(guī)模的高速并行通信方案，即如何解決一個(gè)計(jì)算集群內(nèi)部各個(gè)節(jié)點(diǎn)之間通信速度的瓶頸。

本文的作者提出了一種同步的分布式 SGD 優(yōu)化算法，同時(shí)還提出了幾個(gè)大規(guī)模分布式深度學(xué)習(xí)訓(xùn)練過(guò)程中使用的機(jī)器學(xué)習(xí)方法和優(yōu)化方法，在加速收斂的過(guò)程中保證模型的質(zhì)量沒(méi)有損失。

圖：左圖為 4-chip 的云 TPU v2 設(shè)備，峰值計(jì)算能力為 180 teraFLOPS（每秒萬(wàn)億次浮點(diǎn)運(yùn)算），使用 64GB 的 HBM（高帶寬內(nèi)存）；右圖為使用水冷的 4-chip 云 TPU v3 設(shè)備，峰值計(jì)算能力為 420 teraFLOPS，使用 128GBHBM。TPU v2 設(shè)備可組成最高 256-chip 的計(jì)算集群，稱(chēng)為 TPU Pod，可提供高達(dá) 11.5petaFLOPS 的混合精度吞吐量。TPU v3 Pod 的規(guī)?？蛇_(dá) 1024-chip，是 TPU v2 Pod 的四倍，理論上可提供 107.5 petaFLOPS 的混合精度吞吐量

方法

本文的作者受之前大規(guī)模訓(xùn)練方法的啟發(fā)，在實(shí)驗(yàn)過(guò)程中使用了以下一些技術(shù)：

混合精度：在實(shí)驗(yàn)過(guò)程中，卷積操作使用了 bfloat16 數(shù)據(jù)，這是一種 TPU 上的半精度 16 位浮點(diǎn)數(shù)。此外，卷積層之間的激活函數(shù)也使用了 bfloat16 的格式。為了保證計(jì)算精度與 32 位浮點(diǎn)數(shù)網(wǎng)絡(luò)不相上下的精度，對(duì)于所有的非卷積的操作（例如，批歸一化、損失函數(shù)計(jì)算、梯度求和）都使用了 32 位浮點(diǎn)數(shù)。由于網(wǎng)絡(luò)訓(xùn)練過(guò)程中的主要計(jì)算和內(nèi)存消耗都是在卷積操作上，因此使用 bfloat16 可以獲得更高的訓(xùn)練吞吐量。

學(xué)習(xí)率配置： 先前的一些研究表明，學(xué)習(xí)率應(yīng)當(dāng)與 batch size 成比例。在實(shí)驗(yàn)過(guò)程中，作者使用了線性變化的學(xué)習(xí)率策略進(jìn)行配置（例如，batch size 設(shè)成兩倍，則學(xué)習(xí)率也設(shè)為兩倍）。同時(shí)作者也使用了平緩的學(xué)習(xí)率預(yù)熱（warm-up）方法和學(xué)習(xí)率衰減。

分層自適應(yīng)速率縮放（LARS）： 盡管使用動(dòng)量（momentum）的隨機(jī)梯度下降算法已經(jīng)可以將 batch 最高設(shè)為 8192，但使用 LARS 優(yōu)化器可以達(dá)到 32786 的 batch size 并且對(duì)于模型質(zhì)量沒(méi)有影響。更大的 batch size 也增加了模型在 TPU 集群上執(zhí)行時(shí)的吞吐量。

分布式批歸一化

批歸一化在圖像分類(lèi)任務(wù)中有著不可或缺的作用，它通過(guò)對(duì)一個(gè) mini-batch 內(nèi)的數(shù)據(jù)進(jìn)行歸一化，使得經(jīng)過(guò) batch-norm 層的數(shù)據(jù)服從相同均值與方差的分布，使得下層神經(jīng)元可以更好的對(duì)數(shù)據(jù)分布情況進(jìn)行學(xué)習(xí)。

在分布式訓(xùn)練過(guò)程中，通常讓每個(gè)計(jì)算節(jié)點(diǎn)獨(dú)立的進(jìn)行 batch norm, 這樣的好處是可以大大縮短訓(xùn)練時(shí)間，因?yàn)槊總€(gè)計(jì)算節(jié)點(diǎn)之間無(wú)需額外的通信過(guò)程。在實(shí)驗(yàn)過(guò)程中，作者發(fā)現(xiàn) BN 的批大?。ɡ缬?jì)算節(jié)點(diǎn)的批大小）對(duì)模型的驗(yàn)證準(zhǔn)確率有重要影響。已經(jīng)有研究證明在計(jì)算節(jié)點(diǎn)的批大小小于 32 時(shí)，ResNet-50 的最終訓(xùn)練結(jié)果在驗(yàn)證數(shù)據(jù)上的準(zhǔn)確率并不能收斂。

當(dāng)使用數(shù)據(jù)并行的方法在大規(guī)模計(jì)算機(jī)集群上進(jìn)行部署時(shí)，需要同時(shí)對(duì)全局的 batch size 大小進(jìn)行擴(kuò)大，同時(shí)對(duì)每個(gè)節(jié)點(diǎn)的局部 batch size 進(jìn)行縮小。考慮到 BN 層的影響，作者主要針對(duì)每個(gè)節(jié)點(diǎn)上的 batch size 較小的情況進(jìn)行研究。

作者通過(guò)對(duì)幾個(gè)計(jì)算節(jié)點(diǎn)組成的子節(jié)點(diǎn)做分布式的批歸一化來(lái)實(shí)現(xiàn)對(duì) BN 這一過(guò)程的增強(qiáng)。具體算法如圖所示：

圖：分布式批歸一化算法示意圖，圖中集群包含兩個(gè)計(jì)算節(jié)點(diǎn)

首先各個(gè)節(jié)點(diǎn)計(jì)算獨(dú)立的局部均值與方差

計(jì)算一個(gè)子集群（圖中子集群包含兩個(gè)計(jì)算節(jié)點(diǎn)的）中的分布式均值和方差。

使用分布式均值和方差對(duì)子集群中的所有節(jié)點(diǎn)進(jìn)行歸一化

輸入管道優(yōu)化

訓(xùn)練模型過(guò)程中，輸入管道包括了數(shù)據(jù)讀取、數(shù)據(jù)分析、預(yù)處理、旋轉(zhuǎn)和批量化等操作。如果輸入管道的吞吐量不能和 TPU 等模型管道（前向或反向傳播過(guò)程）的吞吐量相匹配，整個(gè)過(guò)程將會(huì)由于輸入管道的問(wèn)題產(chǎn)生吞吐量上的瓶頸。導(dǎo)致輸入管道與模型管道吞吐量差異的主要原因是專(zhuān)用硬件加速設(shè)備與 CPU 之間的性能差異，因?yàn)槟Ｐ凸艿朗峭耆趯?zhuān)用硬件加速設(shè)備上執(zhí)行的。

在本文中， 作者使用了很多關(guān)鍵的優(yōu)化方法來(lái)解決輸入管道導(dǎo)致的瓶頸。此前，還未有工作對(duì)這些技術(shù)進(jìn)行整合。具體方法如下：

數(shù)據(jù)共享與緩存： 理想情況下，所有的數(shù)據(jù)會(huì)一次性讀取并緩存在內(nèi)存中以備直接使用，但是對(duì)于真實(shí)情況中的大規(guī)模數(shù)據(jù)集這種做法往往是不可行的。由于計(jì)算集群之間是可以共享內(nèi)存與數(shù)據(jù)的，因此在大規(guī)模計(jì)算集群中，作者使用這種數(shù)據(jù)集共享與緩存的方法來(lái)提高輸入管道的吞吐量。

預(yù)提取并計(jì)算： 在計(jì)算當(dāng)前批的數(shù)據(jù)同時(shí)對(duì)下一批的數(shù)據(jù)進(jìn)行提取和處理，當(dāng)前批計(jì)算完時(shí)便可直接提取數(shù)據(jù)使用。

混合 JPEG 解碼與裁剪： 使用原始的編碼數(shù)據(jù)進(jìn)行數(shù)據(jù)增強(qiáng)等操作然后只對(duì)有效的部分進(jìn)行解碼

并行數(shù)據(jù)分析： 對(duì)于輸入管道來(lái)說(shuō)，數(shù)據(jù)分析與處理是非常消耗算力的，多核 CPU 可以使用多線程進(jìn)行加速。

二維梯度求和

本文的作者提出了一種二維梯度求和方法，用于多個(gè)計(jì)算節(jié)點(diǎn)之間的梯度的計(jì)算和傳播。在傳統(tǒng)的一維方法中，梯度求和這一步的時(shí)間復(fù)雜度是 O(n^2)，使用二維求和后，時(shí)間復(fù)雜度可以降到 O(n)。具體計(jì)算方法如下圖所示。

圖：二維環(huán)形梯度傳播，第一階段，藍(lán)色張量在 Y 軸方向進(jìn)行求和，紅色張量在 X 軸方向進(jìn)行求和。第二階段，維度進(jìn)行轉(zhuǎn)換再次求和。

實(shí)驗(yàn)與分析

作者進(jìn)行了多個(gè)實(shí)驗(yàn)，對(duì)文中提到的幾個(gè)技術(shù)細(xì)節(jié)進(jìn)行論證。

分布式批歸一化

分布式歸一化的結(jié)果如下圖所示，實(shí)驗(yàn)使用了 TPU v2 Pod 進(jìn)行訓(xùn)練，并且沒(méi)有使用 LARS 優(yōu)化。

輸入管道優(yōu)化

左圖是逐漸增加每種優(yōu)化方法的實(shí)驗(yàn)結(jié)果，中間的圖是組合優(yōu)化的結(jié)果與逐漸減少其他優(yōu)化方法的結(jié)果對(duì)比，右圖是并行化數(shù)量對(duì)實(shí)驗(yàn)結(jié)果的影響。所有的實(shí)驗(yàn)結(jié)果都以數(shù)據(jù)吞吐量為指標(biāo)。

二維梯度求和

下圖是二維梯度求和算法與一維梯度求和算法的比較，可見(jiàn)使用二維梯度求和在各個(gè)配置的情況下都可以有效的減少分布式求和的時(shí)間。

與已有最好方法的對(duì)比

最后，作者與目前最好的分布式計(jì)算方法進(jìn)行了比較，在準(zhǔn)確率相同的情況下，本文提出的方法相比之前的方法大大減少了時(shí)間消耗。

目前谷歌云已經(jīng)上線 Cloud TPU v3 測(cè)試版，單臺(tái)設(shè)備價(jià)格每小時(shí) 2.4 美元到 8 美元，也不是很貴，你也可以動(dòng)手試試看哦~