24分鐘訓練完ImageNet：ResNet僅需120萬美元，創世界紀錄

24 分鐘訓練ImageNet；僅用120萬美元完成了ResNet-50 的訓練

新智元報道

來源：arXiv

編輯：聞菲，張易

【新智元導讀】你或許還記得Facebook研究人員使用256塊GPU在1小時內訓練完ImageNet的新聞。現在，這項紀錄被刷新：伯克利、TACC和UC戴維斯的研究人員使用新演算法，在24分鐘內訓練完ImageNet的Alex模型，創下新的世界紀錄。同時，研究人員表示他們僅使用了120萬美元的硬體（相比Facebook當時410萬美元），也是使用ResNet-50在一小時內進行ImageNet訓練迄今最低的預算。

此前我們報道過， Facebook 人工智慧實驗室與應用機器學習團隊合作，提出了一種新的方法，能夠大幅加速機器視覺任務的模型訓練過程，僅需 1 小時就訓練完 ImageNet 這樣超大規模的數據集。在最高性能的情況下，Facebook 的新系統，在不犧牲生成模型質量的前提下，每秒使用 256 顆 GPU 訓練 40,000 張圖像。

Facebook 1 小時訓練 ImageNet 的論文發布以後，引起了很多研究者的關注。日前，UC Berkeley、TACC 和 UC Davis 的研究人員發表論文，宣布只用了 24 分鐘就完成了 100-epoch ImageNet 的 AlexNet 訓練；同時，和 Facebook 的研究結果（Goyal et al. 2017）一樣，我們用一小時完成了 90-epoch ImageNet 的 ResNet-50 的訓練。然而，我們的硬體只用了 120 萬美元，比 Facebook 的410 萬低了 3.4 倍。讓我們來關注一下他們的研究成果。

花 24 分鐘用 AlexNet 完成了 100-epoch ImageNet 的訓練；僅用120萬美元完成了 90-epoch ImageNet 的 ResNet-50 的訓練

對於深度學習應用來說，更大的數據集和更大的模型能夠帶來準確度的顯著提升(Amodei et al.2015)，但代價是會花費更長的訓練時間。舉例來說，用英偉達M40 GPU 完成 90-epoch ImageNet-1k 的ResNet-50 訓練要花 14 天。

這一訓練總共要進行 10 的 18 次方的單精度運算。另一方面，世界上目前最快的超級計算機每秒能完成 2*10的17次方 的單精度運算(Dongarra et al. 2017)。如果我們能夠充分利用超級計算機進行 DNN 訓練，理論上我們可以在用5 秒鐘完成 90-epoch ImageNet-1k 的ResNet-50 訓練。然而，目前 DNN 快速訓練的瓶頸在於演算法層面。具體來說，目前的 batch size (e.g. 512) 太小，無法有效利用多處理器。因此，我們需要設計出用於可規模化的 DNN 訓練的並行演算法。

對於大規模 DNN 訓練來說，使用更大的batch size 是必須的。然而，目前方法僅能將 batch size 分別擴展到 1024 （用於 AlexNet）和 8192 （用於 ResNet-50）。更大的 batch size 會帶來明顯的精度損失。本論文中，我們使用了 Layer-wiseAdaptive Rate Scaling（LARS）演算法(You, Gitman, andGinsburg 2017)，使我們可以在不損失精度的條件下將 batch size 擴展到很大（e.g. 32k）。

較大的 batch size 使我們可以有效使用分布系統。我們花了24 分鐘就完成了 100-epoch ImageNet 的 AlexNet 訓練，創造了世界紀錄。和 Facebook 的研究結果（Goyal et al. 2017）一樣，我們用一小時完成了 90-epoch ImageNet 的 ResNet-50 的訓練。然而，我們的硬體只用了 120 萬美元，比 Facebook 的410 萬低了 3.4 倍。

論文摘要

用英偉達 M40 GPU 完成 90-epoch ImageNet-1k 的ResNet-50 訓練要花14 天。這一訓練總共要進行 10 的 18 次方的單精度運算。另一方面，世界上目前最快的超級計算機每秒能完成2*10的17次方的單精度運算(Dongarra et al. 2017)。如果我們能夠充分利用超級計算機進行 DNN 訓練，理論上我們可以在用5 秒鐘完成 90-epoch ResNet-50 訓練。然而，目前 DNN 快速訓練的瓶頸在於演算法層面。具體來說，目前的 batch size 太小(e.g. 512)，無法有效利用多處理器。

在大規模 DNN 訓練中，我們著眼於在固定epochs 中不損失準確度的條件下使用大批量數據並行同步 SGD。 LARS 演算法（You, Gitman, and Ginsburg 2017）使我們可以將 batch size 擴展到很大(e.g. 32K)。我們花 24 分鐘用 AlexNet 完成了 100-epoch ImageNet 的訓練，創造了世界紀錄。和 Facebook 的研究結果（Goyal et al. 2017）一樣，我們用一小時完成了 90-epoch ImageNet 的 ResNet-50 的訓練。然而，我們的硬體只用了 120 萬美元，比 Facebook 的410 萬低了 3.4 倍。

大批量訓練DNN的難題：能夠在5秒內訓練完50層的ResNet嗎？

使用參數伺服器的非同步方法無法保證在大型系統上穩定。對於非常大的DNN訓練，數據並行同步方法更加穩定。這個想法也很簡單——通過為SGD使用大的批量大小，每次迭代的工作可以輕鬆分配到多個處理器。理想情況下，ResNet-50需要772億次單精度操作來處理一個225x225的圖像。如果我們為ImageNet數據集運行90次（epoch），運算數為90*128萬*77.2億（10的18次方）。目前，最強大的超級計算機每秒可完成200×1015次單精度運算。如果有一種演算法可以充分利用超級計算機，我們可以5秒內完成ResNet-50的訓練。

為此，我們需要讓演算法使用更多的處理器，並在每次迭代時載入更多的數據。我們使用一台NVIDIA M40 GPU來說明一台機器的情況。在一定範圍內，更大的批量將使單個GPU的速度更高（圖2）。原因是低級矩陣計算庫將更有效。對於使用ImageNet訓練AlexNet模型，每個GPU的最佳批量大小為512。 如果想要使用許多GPU並使每個GPU都有效，我們需要更大的批量大小。例如，如果我們有16個GPU，那麼我們應該將批量大小設置為16×512 = 8192。理想情況下，如果我們固定總數量訪問量，隨著處理器數量的增加相應呈線性增加批量大小，那麼SGD迭代次數將會線性減小，每次迭代的時間成本保持不變，因此總時間也將隨處理器數量線性減少（表1）。

然而，如果批量大的SGD運行相同數量的epoch，精度通常比小的批量要低得多，而目前也尚未有一種可以有效地使用非常大的批量的演算法。表2顯示了標準基準的目標精度。例如，當設置AlexNet的批量大於1024或ResNet-50的批量大小大於8192時，測試精度將顯著降低（表4和圖3）。

在訓練大的批量的時候，我們需要確保在運行相同數量的epoch的情況下，實現與小的批量差不多的測試精度。在這裡我們固定了epoch的數量，因為：在統計上，一個epoch意味著演算法會觸及整個數據集一次；而在計算上，固定epoch的數量意味著固定浮點運算的數量。訓練大批量最先進的方法包括兩種技術：

線性縮放（Linear Scaling，Krizhevsky 2014）：將批量從B增加到kB，那麼也應該將學習率從η增加到kη。

預熱策略（Warmup Scheme，Goyal等人，2017）：如果使用較大的學習率（η），應該從小的η值開始，然後在前幾個epoch將其增加到大的η。

通過這些技術，研究人員可以在一定範圍內使用相對較大的批量（表3）。然而，我們觀察到，即使最先進的方法，也只能將批量大小縮小到AlexNet的1024，而ResNet-50，更是只能將批量大小縮小為8192。如果我們將AlexNet的批量增加到4096，我們在100個epoch只能達到53.1％的精度（表4）。而我們的目標是，即使使用大批量，也能實現58％的精度。

層次自適應速率縮放（LARS），擴大批量，提升精度

為了提高大批量訓練的準確性，我們使用了一種新的更新學習率（LR）規則。在這裡我們考慮單機情況，使用w = w - η?w來更新權重。使用數據並行方法，可以用相同的方式處理多機器版本。

每個層都有自己的權重w和梯度?w。標準SGD演算法對所有層使用相同的LR（η）。然而，從我們的實驗中，我們觀察到不同的層可能需要不同的LR。原因是||w||2 和 ||?w||2 之間的比率不同的層有很大的不同。

例如，我們觀察到conv1.1層的||w||2/||?w||2僅為20（表5）。但是，對於fc6.1層，||w||2/||?w||2則是3690。為加快DNN訓練，對於fc6.1層，用戶需要使用大的LR。然而，這個大的LR可能導致conv1.1層出現問題。我們認為，這是大批量訓練難以優化的重要原因。Goyal等人在2017年提出了預熱策略來解決這個問題。預熱策略適用於ResNet50訓練（批次大小≤8192）。但是，僅使用這種方法，對於批量大小 >1024的AlexNet和批量大小 >8192的ResNet-50不起作用。

在本文中，我們使用LARS演算法來解決這個問題。基本LR規則在等式（1）中定義。l是縮放因子，我們在AlexNet和ResNet訓練中將l設置為0.001。γ是用戶的調整參數。通常一個好的γ，值都在[1，50]之間。在這個等式中，不同的層可以有不同的LR。實際上，我們向SGD添加動量（用μ表示）和權重衰減（用β表示），並對LARS使用以下序列：

得到每個可學習參數的本地LR，

得到每個層的真實LR，為η=γ×α；

通過?w=?w+βw更新梯度；

通過a =μa+η?w更新加速項a；

用w = w - a來更新權重。

使用這種方法預熱（warmup），用有大的批量的SGD可以實現與基準相同的精度（表6）。為了擴展到更大的批量大小（例如32k），我們需要將本地響應規範化（LRN）更改為批量歸一化（BN）。我們在每個卷積層之後添加BN。從圖3可以清楚地觀察到LARS的影響。

LARS可以幫助ResNet-50保持高的測試精度。當前的方法（線性縮放和預熱）對於批量大小為16k和32k（68％和56％）的精度要低得多。目標精度約為73％。

實驗結果

我們在本文中使用的數據集是ImageNet（1k個類別）（Deng等人，2009）。我們的ResNet-50基線的準確性略低於最先進的結果（73％對76％）。原因是我們沒有使用數據增強。對於沒有數據增強的版本，我們可以得到最先進的精度（90個epoch，73％）。本文中提到的精度意味著Top-1測試精度。時間表示訓練時間。對於KNL實現，我們有兩個版本：

基於Caffe（Jia等人，2014）編寫了KNL代碼，用於單機處理，並使用MPI（Gropp等人，1996）進行KNL集群上不同機器之間的通信。

使用Intel Caffe，支持英特爾MLSL（機器學習擴展庫）的多節點訓練，使用Infiniband作為通信的互連網路。

我們使用TACC Stampede 2超級計算機作為我們的硬體平台。

實現用ImageNet訓練AlexNet當前最快結果：24分鐘100 epoch

此前，英偉達公司報告說，使用一台DGX-1工作站可以在2小時內完成AlexNet在ImageNet訓練90次（epoch）。不過，他們當時使用的是半精度或FP16，計算成本是標準單精度操作的一半。我們用標準的單精度的AlexNet在ImageNet訓練100次。在一個NVIDIA DGX-1站上，批量大小為512需要6小時9分鐘。由於LARS演算法（You，Gitman和Ginsburg 2017），我們能夠使用大的批量大小實現類似的精度。

如果我們將批量大小增加到4096，則在一台NVIDIA DGX-1（一台機器上最快的版本）上只需要2小時10分鐘。因此，在使用LARS的情況下，使用大批量可以顯著加速DNN訓練。

對於批量大小為32K的AlexNet，我們將演算法縮放為512 KNL晶元（約32K處理器或內核）。每個KNL的批量大小是64，所以整體批量大小是32678。我們在24分鐘內完成了100 epoch的訓練。據我們所知，這是使用ImageNet訓練Alex最快的速度，也創下了世界紀錄。總體比較見表7。

當前使用ResNet-50進行ImageNet訓練的最低預算

Facebook（Goyal等人，2017年）在32個CPU和256個NVIDIA P100 GPU（32個DGX-1站）上，在一小時內完成了ResNet-50的90 epoch的ImageNet訓練。考慮到一台DGX-1工作站的費用為129000美元，整個系統的成本約為32×129,000 = 410萬美元。

將批量縮放到32K後，我們可以使用便宜的晶元。我們使用512 KNL晶元，每KNL的批量是64。我們也在一個小時內完成了90 epoch的訓練。但是，我們的系統成本要少得多。我們的KNL晶元的版本是英特爾至強融合處理器7250，成本為2436美元。我們系統的成本只有約2436×512 = 120萬美元。這也是使用ResNet-50在一小時內進行ImageNet訓練迄今最低的預算。

需要指出，我們不隸屬於英特爾或NVIDIA，而且我們對GPU或KNL也沒有任何偏好。我們只想表明，我們可以用更少的預算（便宜3.4倍）實現相同的效果。總體比較見表8。

【號外】新智元正在進行新一輪招聘，飛往智能宇宙的最美飛船，還有N個座位

點擊閱讀原文可查看職位詳情，期待你的加入~

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