今年 6 月，Facebook 實現 1 小時訓練 ImageNet 的成績之後，透過增加批次大小以加快並行訓練速度的方式引發了人們的關注。最近 UC Berkeley 的研究人員們為我們展示了 24 分鐘訓練 ImageNet 的成績，他們將批次大小增加到了 32k。研究人員表示，在同樣成績下，新的方法使用的計算裝置成本（120 萬美元）大大低於 Facebook（410 萬美元）。
對於深度學習應用而言，模型、資料集越大，結果就越精確，但同時也意味著訓練時間的增長。例如，用 ResNet-50 模型在 ImageNet-1k 上完成 90 個 epoch 需要 14 天時間，這樣的訓練共需要 10^18 次單精度計算（single precision operations）。另一方面，世界目前最為強大的超級計算機可以在每秒鐘進行 2×10^17 次單精度計算。如果我們可以完全利用超級計算機的能力，那麼使用 ResNet-50 模型在 ImageNet-1k 上應該可以在五秒內完成 90 個 epoch 的訓練。然而，目前我們面臨的 DNN 訓練瓶頸在演算法層面上。具體來說，目前我們使用的批次過小（如 512），這讓處理器並行處理的效率很低。所以，我們需要為可擴充套件的 DNN 訓練設計有效的並行演算法。
使用大批次SGD是大規模 DNN 訓練的基礎。但是，此前的方法只能將 AlexNet 和 ResNet-50 的批次大小擴大到 1024 和 8192 這樣的程度。更大的批次大小會導致結果準確度的顯著降低。在 UC Berkeley 最近發表的論文中，研究人員使用層級對應的適應率縮放（Layer-wise Adaptive Rate Scaling，LARS）演算法，讓批次大小可以擴大到更大的級別（如 32k）而不損失結果準確度。
大批次可以讓我們更加有效地利用分散式系統的計算能力。UC Berkeley 的研究人員使用 AlexNet 模型在 ImageNet 上訓練 100 個 epoch 只花費了 24 分鐘，這是一個新的世界紀錄。和 Facebook 此前的結果類似（參見：Facebook 新研究：大批次 SGD 準確訓練 ImageNet 僅需 1 小時）新的方法也可以在一小時內使用 ResNet-50 模型在 ImageNet 上訓練 90 個 epoch，但硬體花費約為 120 萬美元——遠少於 Facebook 今年 6 月時使用的 410 萬美元。
圖 1.（a）是資料並行化的一個示例。每個工作器將自己的梯度∇w^j 傳送給控制器，隨後控制器更新自己的權重。隨後，控制器將更新後的權重 w˜傳送給所有工作器。（b）是模型平行化的示例。其中包含一個五層神經網路和本地連線，它被分割成四個部分。只有在跨區域邊緣位置的節點需要在連線處（粗線）共享它們的狀態。即使節點位置與多條邊相連，它也只會對其他部分的節點總共傳送一次自己的狀態。
表 1. 在 ImageNet 資料集上訓練神經網路，tcomp 是計算時間，tcomm 是通訊時間。研究人員將 epoch 數量設為 100。越大的批次大小需要越小的迭代次數。研究人員首先將批次大小設為 512，隨後不斷增加計算裝置數量，由於訓練 ResNet-50 的 Infiniband 網路和 GPU 效率足夠高，單次迭代時間可能接近常數，總訓練時間會不斷減少。
圖 2. 在一定範圍內，大批次提升了系統（GPU）的處理效能。本圖中的資料是 AlexNet 在 ImageNet 資料集上的訓練效果，使用的 GPU 是英偉達 M40，每個 GPU 的批次大小為 512——這是我們能得到的最高速度。Batch=1024 會導致記憶體溢位。
此前，英偉達曾經表示使用 DGX-1 工作站在 ImageNet 中訓練 AlexNet 90 個 epoch 需要花費兩個小時。然而，他們使用了半精度（half-precision）或 FP16，計算資源消耗是標準單精度計算的一半。在本研究中，UC Berkeley 的研究者們使用了標準單精度計算，在英偉達 DGX-1 工作站中，批次大小為 512 時花費了 6 小時 9 分鐘。因為 LARS 演算法的存在，我們可以在保證精度的情況下大幅增加批次大小。如果我們將批次大小增加至 4096，在英偉達 DGX-1 上我們就只需要 2 小時 10 分鐘即可完成任務了（這也是單機最快的成績）。所以，透過 LARS 增加批次大小的方法可以顯著提高 DNN 的訓練速度。在 AlexNet 上，當批次大小為 32k 時，研究人員將演算法擴充套件到 512 個 KNL 晶片（約 32,000 個處理器核心）。每個 KNL 的批次大小為 64，所以總批次大小為 32678。最終，UC Berkeley 實現了 24 分鐘訓練 100 個 epoch 的成績，這也是目前業內最佳的成績。
表 7. 訓練 AlexNet 的速度和硬體花費，批次大小為 32k，研究人員已將 AlexNet 中的區域性響應範數（local response norm）改為批範數。
在今年 6 月份，Facebook 在 ImageNet 上訓練 ResNet-50 時（90 epoch）花費一小時，使用了 32 個 CPU 和 256 個英偉達 P100 GPU（32 個 DGX-1 工作站），目前單個 DGX-1 的價格為 129,000 美元，整個硬體系統的成本為 410 萬美元。
在將批次大小擴充套件到 32k 之後，UC Berkeley 的研究人員們將計算任務交給了相對便宜的計算機晶片——KNL 上，同樣在一個小時的時間裡跑完了 90 個 epoch。他們使用的 KNL 搭載英特爾 Xeon Phi 7250 處理器，單個售價 2436 美元，整個系統花費 120 萬美元。
表 8. 訓練 ResNet-50 的速度和硬體成本，研究中未使用資料增強。
論文：ImageNet Training in 24 Minutes
論文連結：https://www2.eecs.berkeley.edu/Pubs/TechRpts/2017/EECS-2017-155.html
摘要：想要在 ImageNet-1k 上完成 90 個 epoch 的 ResNet-50 模型訓練，使用英偉達 M40GPU 需要花上 14 天。這種訓練共需要 10^18 次單精度計算（single precision operations）。另一方面，世界目前最快的超級計算機可以每秒鐘進行 2×10^17 次單精度計算。如果我們可以充分利用超級計算機的效能用於深度神經網路訊聯，我們應該可以使用 ResNet-50 模型在 5 秒鐘內完成 90 個 epoch 訓練。然而，目前的 DNN 訓練瓶頸在演算法層面上。具體來說，目前我們使用的批次大小（如 512）過於小，無法有效利用多處理器。
對於大規模 DNN 訓練，我們專注於使用大批次資料並行化同步（large-batch data-parallelism synchronous）SGD 在固定 epoch 中保證不損失精度。You、Gitman 與 Ginsburg 在 2017 年提出的 LARS 演算法允許我們將批次大小擴大到 32k 的水平。在此之上，我們只花費了 24 分鐘即完成了在 ImageNet 上，100 個 epoch 的 AlexNet 訓練——這是一個世界紀錄。與 Facebook 的結果相同，我們在 ImageNet 上訓練 ResNet-50 花費了一個小時，而使用的硬體成本僅為 120 萬美元，比 Facebook 使用的 410 萬美元降低了很多。