ImageNet 訓(xùn)練紀(jì)錄再次被刷新！富士通實(shí)驗(yàn)室在 74.7 秒內(nèi)完成了 ImageNet 上訓(xùn)練 ResNet-50 網(wǎng)絡(luò)，使用 2048 個(gè) GPU，準(zhǔn)確率為 75.08%，刷新了此前谷歌 1.8 分鐘的記錄。

在過去兩年中，深度學(xué)習(xí)的速度加速了 30 倍。但是人們還是對(duì) “快速執(zhí)行機(jī)器學(xué)習(xí)算法” 有著強(qiáng)烈的需求。

Large mini-batch 分布式深度學(xué)習(xí)是滿足需求的關(guān)鍵技術(shù)。但是由于難以在不影響準(zhǔn)確性的情況下在大型集群上實(shí)現(xiàn)高可擴(kuò)展性，因此具有較大的挑戰(zhàn)難度。

最近，富士通實(shí)驗(yàn)室的一項(xiàng)研究刷新了一項(xiàng)紀(jì)錄：

論文地址：

https://arxiv.org/pdf/1903.12650.pdf

這項(xiàng)研究在74.7 秒內(nèi)完成了 ImageNet 上訓(xùn)練 ResNet-50 網(wǎng)絡(luò)，使用 2048 個(gè) GPU，準(zhǔn)確率為 75.08%，刷新了此前谷歌 1.8 分鐘的記錄。

表1：ImageNet上訓(xùn)練ResNet-50的記錄

增加 mini-batch 大小，實(shí)現(xiàn)短時(shí)間內(nèi)的高準(zhǔn)確性

基于大數(shù)據(jù)集的深度神經(jīng)網(wǎng)絡(luò) (DNN) 模型在對(duì)象檢測、語言翻譯等領(lǐng)域取得了令人矚目的成果。然而，隨著 DNN 模型和數(shù)據(jù)集規(guī)模的增大，DNN 訓(xùn)練的計(jì)算量也隨之加劇。

具有數(shù)據(jù)并行性的分布式深度學(xué)習(xí)是加速集群訓(xùn)練的一種有效方法。

在這種方法中，集群上啟動(dòng)的所有進(jìn)程都具有相同的 DNN 模型和權(quán)重。每個(gè)過程都用不同的 mini-batch 訓(xùn)練模型，但是來自所有過程的權(quán)重梯度被組合以更新所有權(quán)重。

對(duì)于大型集群，這種通信開銷成為一個(gè)重要的問題。

為了減少大型集群的開銷，該研究增加了 DNN 的 mini-batch 大小，且并行計(jì)算了 DNN 訓(xùn)練。然而，在 minni-batch 訓(xùn)練中，DNN 模型的驗(yàn)證精度普遍較差。

因此，研究者們采用了幾種技術(shù)來增加 mini-batch 的大小，這表明了在迭代中計(jì)算的輸入圖像的數(shù)量，而不會(huì)影響驗(yàn)證的準(zhǔn)確性。

在實(shí)驗(yàn)過程中，本文使用了人工智能橋接云基礎(chǔ)設(shè)備 (AI Bridging Cloud Infrastructure，ABCI) 集群 GPU和自優(yōu)化的 MXNet 深度學(xué)習(xí)框架。并在 ImageNet 上使用 81,920 mini-batch 大小，74.7 秒內(nèi)實(shí)現(xiàn)了 ResNet-50 的 75.08％驗(yàn)證準(zhǔn)確度。

技術(shù)方法三部曲

本文的技術(shù)方法主要分為三個(gè)部分：準(zhǔn)確性改良、框架優(yōu)化和通信優(yōu)化。

A. 準(zhǔn)確性改良

這部分采用了通常用于深度學(xué)習(xí)優(yōu)化器的隨機(jī)梯度下降（SGD）。在對(duì) large mini-batch 進(jìn)行訓(xùn)練時(shí)，SGD 更新的數(shù)量隨著小型批大小的增加而減少，因此提高 large mini-batch 的最終驗(yàn)證精度是一個(gè)很大的挑戰(zhàn)，本文采用了以下技術(shù)。

學(xué)習(xí)速率控制：由于更新數(shù)量較少，需要使用高學(xué)習(xí)率來加速訓(xùn)練。 然而，高學(xué)習(xí)率使得模型訓(xùn)練在早期階段不穩(wěn)定。 因此，我們通過使用逐漸提高學(xué)習(xí)率的預(yù)熱 (warmup) 來穩(wěn)定 SGD。此外，對(duì)于某些層，所有層的學(xué)習(xí)速率都太高了，還通過使用層次自適應(yīng)速率縮放（LARS）來穩(wěn)定訓(xùn)練，LARS 根據(jù)規(guī)范權(quán)重和梯度調(diào)整每層的學(xué)習(xí)速率。

其它技術(shù)：據(jù)報(bào)道，標(biāo)簽平滑提高了 32,768 個(gè) mini-batch 的準(zhǔn)確性。本文也采用了這種方法，并對(duì) 81920 個(gè) mini-batch 進(jìn)行了精度改進(jìn)。

batch 標(biāo)準(zhǔn)化層的均值和方差的移動(dòng)平均 (moving average) 在每個(gè)過程中獨(dú)立計(jì)算，而權(quán)重是同步的。這些值在 large mini-batch 上變得不準(zhǔn)確；因此，本文調(diào)整了一些超參數(shù)來優(yōu)化移動(dòng)平均線。

B. 框架優(yōu)化

我們使用了 MXNet，MXNet 具有靈活性和可擴(kuò)展性，能夠在集群上高效地訓(xùn)練模型。然而，在中小型集群環(huán)境中只占總時(shí)間的一小部分的處理方式可能成為大規(guī)模集群環(huán)境中的瓶頸。我們使用了幾個(gè)分析器來分析 CPU 和 GPU 性能，找出了瓶頸。我們對(duì)瓶頸進(jìn)行了優(yōu)化，提高了訓(xùn)練吞吐量。

1) 并行 DNN 模型初始化：

在數(shù)據(jù)并行分布式深度學(xué)習(xí)中，必須初始化所有層，使所有進(jìn)程的權(quán)重相同。通常，根進(jìn)程初始化模型的所有權(quán)重。然后，進(jìn)程將這些權(quán)重傳遞 (broadcast) 給所有進(jìn)程。傳遞時(shí)間隨著進(jìn)程數(shù)量的增加而增加，在有成千上萬個(gè)進(jìn)程進(jìn)行分布式深度學(xué)習(xí)時(shí)，其成本不可忽視。

因此，我們采用了其他初始化方法，即每個(gè)進(jìn)程具有相同的種子并并行地初始化權(quán)重。這種方法無需 broadcast 操作就可以同步初始權(quán)重。

2) GPU 上的 Batch Norm 計(jì)算：

每層的 norm 計(jì)算都需要使用 LARS 更新權(quán)重。與 GPU 上的內(nèi)核數(shù)量相比，ResNet-50 的大多數(shù)層沒有足夠的權(quán)重。如果我們在 GPU 上計(jì)算每一層的 weight norm，線程數(shù)不足以占據(jù)所有 CUDA 核心。因此，我們實(shí)現(xiàn)了一個(gè)特殊的 GPU 內(nèi)核，用于 batched norm 計(jì)算到 MXNet。該 GPU 內(nèi)核可以啟動(dòng)足夠數(shù)量的線程，并且可以并行計(jì)算層的范數(shù)。

C. 通信優(yōu)化

分布式并行深度學(xué)習(xí)要求所有 reduce 通信在所有進(jìn)程之間交換每一層的梯度。在大集群環(huán)境中，由于每個(gè) GPU 的 batch size 較小，使得通信時(shí)間變長，計(jì)算時(shí)間變短，因此 reduce communication 開銷是不可忽略的。為了克服這些問題，我們采用了以下兩種優(yōu)化方法。

1) 調(diào)整通信的數(shù)據(jù)大小

2) 通信的優(yōu)化調(diào)度

實(shí)驗(yàn)設(shè)置與實(shí)驗(yàn)結(jié)果

我們使用 ABCI 集群來評(píng)估基于 MXNet 的優(yōu)化框架的性能。ABCI 集群的每個(gè)節(jié)點(diǎn)由兩個(gè) Xeon Gold 6148 CPU 和四個(gè) NVIDIA Tesla V100 SXM2 GPU 組成。此外，節(jié)點(diǎn)上的 GPU 由 NVLink 連接，節(jié)點(diǎn)也有兩個(gè) InfiniBand 網(wǎng)絡(luò)接口卡。圖 1 為 ABCI 集群節(jié)點(diǎn)結(jié)構(gòu)示意圖。

圖 1：ABCI 集群中一個(gè)計(jì)算節(jié)點(diǎn)的示意圖。它由兩個(gè) GPU、四個(gè) GPU 和兩個(gè)連接到相應(yīng) CPU 的 HCA 組成。

我們使用混合精度方法，使用半精度浮點(diǎn)數(shù)計(jì)算和通信，并使用單精度浮點(diǎn)數(shù)更新權(quán)重。我們使用了原始優(yōu)化器，它可以很好地控制學(xué)習(xí)率。除了穩(wěn)定訓(xùn)練精度外，我們還使用了 warmup 和 LARS 技術(shù)。

我們對(duì) ResNet-50 訓(xùn)練的測量依據(jù) MLPerf v0.5.0 規(guī)則。也就是說，我們度量了從 “run start” 到 “run final” 的運(yùn)行時(shí)間，其中包括初始化和內(nèi)存分配時(shí)間。

結(jié)果表明，優(yōu)化后的 DNN 框架在 74.7 秒內(nèi)完成了 ImageNet 上 ResNet-50 的訓(xùn)練，驗(yàn)證精度為 75.08%。

圖 2：優(yōu)化后的框架的可擴(kuò)展性用實(shí)線表示，虛線表示理想曲線。

我們還測量了 ResNet-50 的可擴(kuò)展性。圖 2 顯示了根據(jù) GPU 數(shù)量計(jì)算的吞吐量。在圖 2 中，虛線表示理想的每秒圖像吞吐量，實(shí)線表示我們的結(jié)果。如圖表明，直到 2048 個(gè) GPU，我們的框架的可擴(kuò)展性都非常好。使用 2048 個(gè) GPU 的吞吐量為每秒 170 萬張圖像，可擴(kuò)展性為 77.0%。

圖 3：在 49152 個(gè)或更大的 mini-batch 訓(xùn)練中，top-1 驗(yàn)證精度的變化

圖 3 顯示了 81,920 個(gè)或更大的 mini-batch 訓(xùn)練中 top-1 驗(yàn)證精度的結(jié)果。從圖 3 中可以看出， mini-batches 超過 81,920 個(gè)的驗(yàn)證精度低于 74.9%，不符合 MLPerf 規(guī)定。因此，ImageNet 數(shù)據(jù)集一個(gè) epoch 的圖像數(shù)量為 1,280,000 張，如果使用 81,920 mini-batch，一個(gè) epoch 中的更新數(shù)量僅為 16 張，其中更新總數(shù)為 1,440 張。這個(gè)數(shù)字太小，SGD 求解器無法訓(xùn)練 DNN 權(quán)重。因此，使用大的 mini-batch 是一個(gè)很大的挑戰(zhàn)，我們嘗試使用盡可能大的 mini-batch。

如表 1 所示，與其他工作相比，81,920 mini-batch size 已經(jīng)很大，驗(yàn)證精度達(dá)到 75% 以上。

圖 4：訓(xùn)練精度與驗(yàn)證精度之比較

圖 4 顯示了訓(xùn)練精度與驗(yàn)證精度的對(duì)比。從圖中可以看出，使用 batch normalization 和 label smoothing 技術(shù)，我們的驗(yàn)證精度結(jié)果并沒有過擬合。

結(jié)論

我們開發(fā)了一種新的技術(shù)，可以在大規(guī)模 GPU 集群上使用 large mini-batch，而不會(huì)降低驗(yàn)證精度。我們將該技術(shù)應(yīng)用到基于 MXNet 的深度學(xué)習(xí)框架中。使用 81920 minibatch size，我們的 DNN 訓(xùn)練結(jié)果在 74.7 秒內(nèi)訓(xùn)練完 ResNet-50，驗(yàn)證精度達(dá)到 75.08%。