Google AI 博客：通過數據回傳加速神經網絡訓練

過去十年來，神經網絡的訓練速度得到了大幅提高，使得深度學習技術在許多重要問題上的應用成爲可能。隨着摩爾定律即將走向終結，通用處理器的的改進未取得明顯成效，機器學習社區越來越多地轉向專用硬件來謀求額外的加速。

GPU 和 TPU 針對高度並行化的矩陣運算進行了優化，而矩陣運算正是神經網絡訓練算法的核心組成部分。在高層次上，這些加速器可以通過兩種方式來加速訓練。首先，它們可以並行處理更多的訓練樣本；其次，它們可以更快地處理每個訓練樣本。我們知道，並行處理更多訓練樣本的加速是有限制的，但構建更快的加速器，還能繼續加速訓練嗎？

遺憾的是，並非所有訓練管道中的操作都是在加速器上運行的，因此，不能簡單地依靠速度更快的加速器來繼續推動訓練速度的提升。例如，訓練管道中的早期階段（如磁盤 I/O 和數據預處理）涉及到的操作並不能從 GPU 和 TPU 中受益。隨着加速器的改進超過 CPU 和磁盤的改進，這些早期階段日益成爲瓶頸，從而浪費加速器的容量，並限制了訓練速度的提升。

圖爲代表了許多大型計算機視覺的訓練管道示例。在應用小批量隨機梯度下降（mini-batch stochastic gradient descent，（SGD））更新之前的階段一般不會受益於專用硬件加速器。

考慮這樣一個場景：代碼從上游到加速器的時間是代碼在加速器上運行時間的兩倍，對於目前一些工作負載來說，這個場景已經很現實了。即使代碼以流水線的方式並行執行上游和下游階段，上游階段的代碼也會佔據大量的訓練時間，加速器將有 50% 的時間處於閒置狀態。在這種情況下，構建更快的加速器來加速訓練根本無濟於事。也許，可以通過增加工作量和額外的計算資源來加快輸入管道的速度。但這種努力很費時，並且還偏離了提高預測性能的主要目標。對於非常小的數據集，人們可以離線預計算擴展後的數據集，並在內存中加載整個經過預處理後的數據集，但這一做法對於大多數機器學習訓練場景來說，行不通。

在論文《通過數據回傳加速神經網絡訓練》（Faster Neural Network Training with Data Echoing）中，我們提出了一種簡單的技術，可以重用（或“回傳”）早期管道階段的中間輸出，以回收閒置的加速器容量。與其等待更多數據可用，倒不如我們只需利用已有的數據來保持加速器的繁忙狀態。

左圖：在沒有數據回傳的情況下，下游算力有 50% 的時間處於閒置狀態。右圖：回傳因子爲 2 的數據回傳，回收了下游算力。

重複數據以加快訓練

假設有這樣一種情況，對一批訓練數據進行讀取和預處理所需的時間，是對該批數據執行單個優化步驟所需時間的兩倍。在這種情況下，在對已預處理的批處理進行第一個優化步驟之後，我們就可以重用該批處理，並在下一個批處理準備好之前執行第二個步驟。在最好的情況下，重複的數據和新數據一樣有用，我們將會看到訓練的速度提高了兩倍。不過，在現實中，由於重複數據不如新數據有用，數據回傳帶來的速度提升略小，但與讓加速器處於閒置狀態相比，它仍然可以帶來顯著的速度提升。

在給定的神經網絡訓練管道中，通常有幾種方法可以實現數據回傳。我們提出的技術涉及到將數據複製到訓練管道中某個位置的數據洗牌緩衝區（shuffle buffer）中，但是我們可以在給定管道中產生瓶頸的任何階段之後自由地插入這個緩衝區。當我們在批處理前插入緩衝區時，我們稱之爲技術樣本回傳（technique example echoing）；而當我們在批處理之後插入緩衝區時，我們稱之爲技術批處理回傳（technique batch echoing）。樣本回傳會在樣本級別上對數據進行洗牌，而批處理回傳則對重複批的序列進行洗牌。我們還可以在數據擴展之前插入緩衝區，這樣，每一份重複數據的副本都略有不同（因此更接近於新鮮樣本）。在不同階段之間放置數據洗牌緩衝區的不同版本的數據回傳中，所能提供最大速度提升的版本取決於具體的訓練管道。

跨工具負載的數據回傳

那麼，重用數據能有多大用呢？我們在五個神經網絡訓練管道上嘗試了數據回傳，這些訓練管道涵蓋了三個不同的任務：圖像分類、語言建模和目標檢測，並測量了達到特定性能目標所需的新鮮樣本的數量。在超參數調優過程中，我們選擇的目標與基線能夠可靠達到的最佳結果相匹配。我們發現，數據回傳讓我們能夠用更少的新鮮樣本達到目標性能，這表明，重用數據對於降低各種任務的磁盤 I/O 非常有用。在某些情況下，重複數據幾乎和新鮮數據一樣有用：在下圖中，擴展前的樣本回傳，幾乎按重複因子減少了所需的新鮮樣本的數量。

當每個數據項重複兩次時，數據回傳會減少或不改變達到目標樣本外性能所需新鮮樣本的數量。虛線表示我們所期望的值，假設重複樣本和新樣本一樣有用。

縮短訓練時間

當加速器上游的計算佔用訓練時間時，數據回傳可以加速訓練。我們測量了在訓練管道中實現的訓練提速，該管道由於來自雲存儲的流式傳輸訓練數據的輸入延遲而成爲瓶頸，這對於當今許多大規模生產工作負載或任何人來說，通過網絡從遠程存儲系統流式傳輸訓練數據都是現實的。我們在 ImageNet 數據集上訓練了 ResNet-50 模型，發現數據回傳提供了顯著的訓練提速，在這種情況下，使用數據回傳時，速度要快三倍以上。

數據回傳可以減少 ResNet-50 在 ImageNet 上的訓練時間。在此實驗中，從雲存儲中讀取一批訓練數據，比使用每批數據進行訓練步驟的代碼要長 6 倍。圖例中的回傳因子指的是每個數據項重複的次數。虛線表示如果重複樣本與新樣本一樣有用且不存在回傳開銷時的期望值。

數據回傳保持了預測性能

儘管人們可能會擔心重用數據會損害模型的最終性能，但我們發現，在我們測試過的任何工作負載中，數據回傳並沒有降低最終模型的質量。

通過比較在訓練期間獲得最佳樣本外性能的單個實驗，無論有無數據回傳，都表明重用數據並不會損害最終模型質量。此處的驗證交叉熵（cross entropy）相當於對數困惑度（log perplexity）。

隨着 GPU 和 TPU 等專用加速器的改進繼續超過通用處理器，我們預計，數據回傳和類似策略將會成爲神經網絡訓練工具包中越來越重要的部分。

作者介紹：

Dami Choi，Google AI 學生研究員。George Dahl，高級研究科學家，供職於 Google Research。

原文鏈接：

https://ai.googleblog.com/2020/05/speeding-up-neural-network-training.html