繼1小時訓練ImageNet之後，批量訓練擴展到3萬2千樣本

選自arXiv

機器之心編譯

參與：蔣思源、李亞洲、路雪

自 Facebook 發布 1 小時訓練 ImageNet 論文以來，很多研究者都在關注如何使用並行訓練來提高深度學習的訓練速度。Facebook 提出了一種提高批量大小的分散式同步 SGD 訓練方法，而 Yang You 等人在 Facebook 的論文上更進一步採用層級對應的適應率縮放（LARS）來對每一層網路使用不同的學習率。他們在 AlexNet 和 ResNet-50 模型上分別實現了 8129 和 32768 的批量大小，而且在加速訓練的同時還保證了模型的準確度。

論文地址：https://arxiv.org/abs/1708.03888

加速訓練大型神經網路最自然的方式就是在多塊 GPU 上使用數據平行化。為了將基於隨機梯度的方法擴展到更多的處理器，我們需要增加批量大小以充分利用每塊 GPU 的計算力。然而，在提高批量大小的情況下保持住神經網路的準確度就顯得十分重要了。目前，最優秀的方法是與批量大小成正比地提高學習率（Learning Rate /LR），並使用帶有「warm-up」策略的專用學習率來克服優化困難。

通過在訓練過程中控制學習率（LR），我們可以在 ImageNet-1K 訓練中高效地使用大批量梯度下降。例如，在 AlexNet 中應用 Batch-1024 和在 ResNet-50 中應用 Batch-8192 都取得了不錯的效果。然而對於 ImageNet-1K 訓練來說，最優秀的 AlexNet 只能擴展到 1024 的批量大小，最優秀的 ResNet-50 只能擴展到 8192 的批量大小。我們不能將學習率擴展到一個較大的值。為了將大批量訓練擴展到一般網路或數據集，我們提出了層級對應的適應率縮放（Layer-wise Adaptive Rate Scaling/LARS）。LARS LR 基於權重的範數和梯度的範數在不同的神經網路層級中使用不同的學習率。通過使用 LARS LR，我們將 ResNet-50 的批量大小擴展到 32768，將 AlexNet 的批量大小擴展到 8192。大批量同樣能完全利用系統的計算能力。例如，AlexNet 模型在 DGX-1 工作站（8 P100 GPU）上使用 Batch-4096 要比 Batch-512 訓練 ImageNet 快 3 倍。

1 前言

使用 AlexNet 模型（2）在 NVIDIA K20 GPU 上訓練 ImageNet 數據集（1）需要 6 天才能實現 58% 的 top-1 準確度（3）。因此，擴展和加速 DNN 的訓練對於應用深度學習來說極其重要。

我們關注數據並行化的小批量隨機梯度下降訓練（4），該演算法在許多如 Caffe（5）和 TensorFlow（6）那樣的流行深度學習框架中都是頂尖的優化方法。我們在該研究中使用的是英偉達 GPU。為了加速 DNN 的訓練，我們需要將演算法擴展到更多的處理器中。所以為了將數據平行化的 SGD 方法擴展到更多的處理器中，我們需要增加批量大小。增加 GPU 數量的同時增加批量大小，並能夠使每一塊 GPU 的工作保持穩定，並能令每塊 GPU 的計算資源都得到充分利用。

為了更進一步提升大批量 AlexNet 的測試準確度，並令大批量訓練應用到一般的神經網路和數據集中，我們提出了層級對應的適應率縮放（Layer-wise Adaptive Rate Scaling/LARS）。LARS 基於權重的範數（||w||）和梯度的範數（||?w||）對不同的層級使用不同的學習率。LARS 背後的原因是我們觀察到權重範數和梯度範數的比率（||w||/||?w||）在每一層的區別非常大。

我們在 AlexNet 模型上的批量由 128 提高到 8129，但實現了相似的準確度。而對於 ResNet-50 來說，我們在 ImageNet 訓練中成功地將批量大小提高到 32768。

2 背景和相關研究

2.1 小批量隨機梯度下降的數據並行化

現在令 w 代表 DNN 的權重、X 代表訓練數據、n 為 X 中的樣本數，而 Y 代表訓練數據 X 的標註。我們令 x_i 為 X 的樣本，（x_i，w）為 x_i 和其標註 y_i（i ∈ {1, 2, ..., n)）所計算出的損失。一般來說，我們使用如交叉熵函數那樣的損失函數。DNN 訓練的目標是最小化方程（1）中的損失函數。

在第 t 次迭代中，我們使用前向和反向傳播以求得損失函數對權重的梯度。然後，我們使用這個梯度來更新權重，根據梯度更新權重的方程（2）如下：

其中 η 為學習率。該方法被稱作隨機梯度下降（SGD）。通常，我們並不會只使用一個樣本計算損失和梯度，我們每次迭代會使用一個批量的樣本更新權重。現在，我們令第 t 次迭代的批量大小為 B_t，且 B_t 的大小為 b。然後我們就可以基於以下方程（3）更新權重：

這種方法叫作小批量隨機梯度下降（Mini-Batch SGD）。為了簡化表達方式，我們可以說方程（4）中的更新規則代表我們使用權重的梯度?w_t 更新權重 w_t 為 w_t+1。

2.2 大批量訓練的困難

2.3 學習率（LR）

2.4 頂尖的大批量訓練

3 ImageNet-1k 訓練

3.1 重現和延伸 Facebook 的研究結果

與 facebook 的論文類似，我們使用預熱策略（warm-up policy）和線性縮放策略來獲取學習率。但是我們和 Facebook 有兩點不同：（1）我們將學習率提升得更高，Batch-256 使用的基本學習率是 0.2，而 Batch-8192 使用的基本學習率是 6.4。（2）我們使用多個規則（poly rule）而不是多步規則（multistep rule）來更新學習率。

• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour（Facebook論文）：https://arxiv.org/abs/1706.02677

圖 1：從這張圖片中，我們可以清楚地看到我們使用 8k 批量大小的準確度和 256 批量大小的準確度在 90 個 epoch 中達到同等水平。該實驗在一個 NVIDIA DGX-1 工作站上完成。

表 1：ResNet-50 模型在 ImageNet 數據集上使用多個學習率規則的情況。

3.2 使用 AlexNet 訓練 ImageNet

3.2.1 學習率的線性縮放和預熱策略

我們使用批量為 512 的 BVLC-AlexNet 作為我們的基線，其在 100 個 epoch 中實現了 0.58 的準確度。在該實驗中，我們使用多個學習率規則。基本的學習率是 0.01，而多個規則的學習率為 2。我們的目標是使用 Batch-4096 和 Batch-8192 在 100 個 epoch 實現 0.58 的準確度。

表 2：AlexNet 模型使用多個學習率策略在 ImageNet 數據集中的訓練情況。我們使用學習率的線性縮放和預熱策略，這些是 Facebook 論文（7）中使用的主要技術。

3.2.2 用於大批量訓練的批歸一化（BN）

為在 AlexNet 上實現大批量訓練，我們嘗試了許多不同技術（如，數據清洗、數據縮放（data scaling）、多步 LR、最小 LR 調參）。我們觀察到只有批歸一化（BN）能夠提升準確度。

3.2.3 用於大批量訓練的層級對應的適應率縮放（LARS）

為了提高大批量 AlexNet 的準確度，我們設計了一個更新學習率的新規則。如前文所述，我們使用 w = w ? η?w 來更新權重。每一層都有自己的權重 w 和梯度 ?w。標準 SGD 演算法在所有層上使用相同的 LR（η）。但是，從實驗中我們觀察到不同的層可能需要不同的學習率。

表 3：AlexNet-BN 使用多個 LR 策略在 ImageNet 數據集上的訓練情況。

圖 2：從上圖中，我們清楚地觀察到在添加批歸一化後，泛化性能變得更好。

表 5：AlexNet_BN 使用 LARS 學習率和多個規則在 ImageNet 數據集上的訓練情況。

表 6：AlexNet 模型在 ImageNet 數據集中的訓練情況

圖 3：所有例子只運行 100 個 epoch。我們能明顯觀察到 LARS 學習率的影響。

4 實驗結果

4.1 系統和庫

4.2 實現細節

4.3 訓練 ResNet 50 的頂級結果

通過使用 LARS，我們能把 ResNet50 模型的 batch 大小擴展到 32K。我們使用 Batch-256 作為基線。在最高準確度上，8K 的 Batch 損失了 0.3% 的準確度，32K 的 Batch 損失了 0.7% 的準確度。我們相信這些準確度損失能夠通過超參數調整彌補。我們的基線要比頂級的結果（7）（17）略低，因為我們並未使用數據增強。

圖 4：通過使用 LARS，我們能把 ResNet-50 的批量大小擴展到 32K

表 7：ResNet-50 模型在 ImageNet 數據集上的訓練情況

表 8:AlexNet-BN 的速度與時間

4.4 使用大批量的收益

圖 5：在 DGX-half（4 塊 GPU）上，512 的 Batch 與 4096 的 Batch 有同樣的速度。在 DGX（8 塊 GPU）上，4096 的 Batch 是 512 Batch 速度的 2.6 倍。因此，大批量在計算能力增強的時候會有收益。

表 10：AlexNet 的速度和時間

5 結論

優化難題導致大批量訓練的準確度損失。只使用線性縮放和預熱策略（warmup scheme）這樣的方法對複雜的應用而言並不足夠，例如使用 AlexNet 來訓練 ImageNet。我們提出層級對應的適應率縮放（Layer-wise Adaptive Rate Scaling/LARS），它基於權重的範數和梯度的範數在不同層級上使用不同的學習率。在實驗中，LARS 表現出了極高的效率。通過使用 LARS，把進行 ImageNet 訓練的 AlexNet 模型的批量大小從 128 增加到 8192 的時候，我們依然能得到同樣的準確度。我們也能把 ResNet-50 的批量大小擴展到 32768，大批量也能充分使用系統的計算能力。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！