Training Deep Nets with Sublinear Memory Cost

《Training Deep Nets with Sublinear Memory Cost》筆記

摘要

我們提出了一種減少深度神經網絡訓練時記憶體消耗的系統性方法。具體來說，我們設計了一個算法，訓練一個 n n 層網絡僅耗費 O(n−−√)O(n) 的記憶體，每個mini-batch隻需要一個額外的前向計算成本。由于許多最先進的模型已經達到了GPU顯存的上限，我們的算法允許探索更深入更複雜的模型，并有助于推進深度學習研究的創新。我們專注于降低訓練期間存儲中間特征圖和梯度的記憶體成本。計算圖分析用于自動原地操作和記憶體共享優化。我們證明通過一種記憶體利用更高效的訓練算法以及一點額外的計算成本，可以通過計算來換取記憶體資源。在極端情況下，我們的分析還表明，記憶體消耗可以減少到 O(logn) O ( log ⁡ n ) ，而正向計算的額外開銷隻需 O(nlogn) O ( n log ⁡ n ) 。我們的實驗表明，在ImageNet上我們可以将1000層深度殘差網絡的記憶體成本從48G降低到7G。類似地，在非常長的序列中訓練複雜遞歸神經網絡時亦可顯着的降低記憶體成本。

介紹

在本文中，我們提出了一種減少深度神經網絡訓練時記憶體消耗的系統性方法。我們主要關注于減少存儲中間結果（特征映射）和梯度的記憶體成本，因為在許多常見深度架構中，與中間特征映射的大小相比，參數的大小相對較小。我們使用計算圖分析來執行自動原地操作和記憶體共享優化。更重要的是，我們提出了一種新的以計算交換記憶體的方法。進而，我們給出了一個實用的算法，其花費 O(logn) O ( log ⁡ n ) 的記憶體用于特征映射，僅以雙倍正向計算成本訓練n層網絡。有趣的是，我們還表明，在極端情況下，可以使用 O(logn) O ( log ⁡ n ) 的記憶體存儲特征映射來訓練一個 n n 層網絡。

我們最近見證了深度神經網絡在許多領域的成功[8]，如計算機視覺、語音識别、自然語言處理和強化學習。深度神經網絡新架構方面的創新帶來了許多成功。卷積神經網絡[15,14,13,10]對空間模式進行模組化并給出計算機視覺任務中的頂尖結果。遞歸神經網絡，如長期短期記憶[12]，在序列模組化和結構預測中顯示出鼓舞人心的結果。這些新模型的一個共同趨勢是使用更深的架構[18,14,13,10]來捕捉大量訓練資料中的複雜模式。由于存儲特征映射及其梯度的成本随網絡深度線性變化，是以我們探索深層模型的能力受到裝置（通常為GPU）記憶體的限制。例如，如[11]所述，目前最先進的一個模型已經耗盡了記憶體。從長遠來看，理想的機器學習系統應該能夠從越來越多的訓練資料中不斷學習。由于最佳模型大小和複雜性通常随着訓練資料增多而增加，是以節約記憶體的訓練算法是非常重要的。

減少記憶體消耗不僅使我們能夠訓練更大的模型。它還可以實作更大的批量大小，進而提高批量操作符的裝置使用率和穩定性，如批量歸一化[13]。對于記憶體有限的裝置，它有助于改善記憶體局部性，并可能産生更好的記憶體通路模式。它還使我們訓練深度卷積神經網絡時能夠從模型并行性切換到資料并行性，這在某些情況下可能是有益的。我們的解決方案使我們能夠訓練更深的卷積神經網絡，以及具有更大展開步長的遞歸神經網絡。結合本文提出的記憶體優化技術，我們為深度學習架構提供指導。我們還将公開我們記憶體優化算法的實作。

基于計算圖的記憶體優化

我們首先回顧一下計算圖的概念和記憶體優化技術。其中一些技術已經被現有的架構如Theano[5,4]，Tensorflow[1]和MXNet[6]所使用。計算圖由表示操作之間依賴關系的操作節點和邊組成。圖1給出了一個雙層全連接配接神經網絡計算圖的例子。這裡我們使用粗粒度的向前和向後操作來使圖更簡單。我們通過隐藏權重節點和權重的梯度來進一步簡化圖。在實踐中使用的計算圖可能更複雜，并且混合包含粗/細粒度操作。本文提出的分析可以直接用于那些更一般的情況。

圖1 兩層全連接配接神經網絡訓練過程的計算圖和可能的記憶體配置設定計劃。每個節點代表一個操作，每個邊表示操作之間的依賴關系。具有相同顔色的節點共享記憶體以存儲輸出或反向傳播的梯度。為使圖更清楚，我們省略了圖中權重及其輸出梯度節點，并假定在後向運算時會計算權重的梯度。我們還注釋了使用原地和共享政策的兩個地方。

一旦給出網絡配置（前向圖），我們就可以建構相應的後向通道進行梯度計算。可以通過以逆向拓撲順序周遊配置來建構後向路徑，并且像在正常反向傳播算法中那樣應用後向運算符。圖1中的後向路徑顯式給出了梯度計算步驟，使得訓練中的梯度計算步驟簡化為整個計算圖表（包括梯度計算路徑）上的正向通過。顯式梯度路徑還提供了一些其他好處（例如，能夠計算更高階梯度），這超出了本文所讨論的範圍。

在訓練深度卷積/循環網絡時，通常使用大部分記憶體來存儲中間輸出和梯度。這些中間結果中的每一個都對應于圖中的一個節點。智能配置設定算法能夠通過盡可能記憶體共享來為這些節點配置設定最少量的記憶體。圖1展示了例子中兩層神經網絡可能的配置設定計劃。可以使用兩種類型的記憶體優化

• 原地操作：直接将輸出值存儲到輸入值的記憶體中。
• 記憶體共享：回收不再需要的中間結果使用的記憶體并用于另一個節點。

圖1中的配置設定計劃包含兩種情況的例子。第一個sigmoid轉換使用原地操作來節省記憶體，而其後向操作又進行了重用。softmax梯度可以與第一個全連接配接層的梯度共享記憶體。特定情況下應用這些優化可能會導緻錯誤。例如，如果一個操作的輸入仍然是另一個操作所需要的，那麼對輸入進行原地操作将導緻錯誤的結果。

我們隻能在生命周期不重疊的節點之間共享記憶體。有多種方法可以解決這個問題。一個選擇是以每個變量作為節點，變量之間的重疊生命周期為邊，構造沖突圖，然後運作圖着色算法。這将花費 O(n2)O(n2) 的計算時間。我們采用更簡單的啟發式方法，僅需 O(n) O ( n ) 的時間。該算法如圖2所示。它以拓撲順序周遊圖，并使用計數器來訓示每條記錄的活性。當沒有其他操作等待其輸入時，可能會産生原地操作。另一個節點使用回收标簽時會進行記憶體共享。這也可以作為圖周遊的動态運作時算法，并使用垃圾回收器來回收過時的記憶體。我們将其用作靜态記憶體配置設定算法，在執行開始之前為每個節點配置設定記憶體，以避免在運作時垃圾收集的開銷。

圖2 計算圖上的記憶體配置設定算法。每個節點都關聯一個活性計數器，用于記錄要填充的操作。時間标記用于訓示記憶體共享。如果目前操作是剩下的唯一操作（輸入的計數器等于1），則可以執行原地操作。當節點計數器變為零時，可以回收節點的标簽。

深度學習架構指南 正如我們從圖2中的算法示範圖中所看到的那樣。資料依賴性導緻每個輸出的使用期更長，并增加了大型網絡的記憶體消耗。對于深度學習架構來說非常重要是

• 以最小方式聲明梯度運算符的依賴性要求。
• 對依賴性資訊應用活性分析并啟用記憶體共享。

聲明最小的依賴關系很重要。例如，如果sigmoid-backward也取決于第一個fullc-forward的輸出，則圖1中的配置設定計劃将不可能實作。依賴性分析通常可以将 n n 層網絡的深度網絡預測的記憶體占用從 O(n)O(n) 減少到接近 O(1) O ( 1 ) ，因為可以在每個中間結果之間進行共享。該技術也有助于減少訓練的記憶體占用，盡管隻是一個常數因素。

以計算換記憶體

通用方法

上一節介紹的技術可以減少深度神經網絡訓練和預測的記憶體占用。然而，由于大多數梯度算子都依賴于正向通路的中間結果，是以訓練 n n 層卷積網絡或序列長度為 nn 的遞歸神經網絡，我們仍需要 O(n) O ( n ) 的記憶體用于中間結果。為了進一步減少記憶體，我們建議删除一些中間結果，并在需要時借助額外的前向計算恢複它們。

更具體地說，在反向傳播階段，我們可以通過從最近的記錄結果運作前向計算來獲得丢棄的中間結果。為了更清楚地呈現這個想法，我們在Alg. 1 中展示了一個直鍊前饋神經網絡的簡化算法。具體來說，神經網絡被分為幾個部分。該算法隻記住每段的輸出并删除段内的所有中間結果。反向傳播期間在段内重新計算丢棄的結果。是以，我們的記憶體開銷僅需覆寫所有段的輸出以及逐段反向傳播時的最大記憶體占用。

隻要我們可以将圖分成幾部分，Alg. 1 也可以推廣到其他常用的計算圖。但是，直接應用ALG. 1有兩個缺點：

• 使用者必須手動劃分圖并編寫定制的訓練循環;
• 我們無法從上節中介紹的其他記憶體優化中受益。

我們通過引入基于相同思想的通用梯度圖構造算法來解決這一問題。該算法在Alg. 2中給出。在該算法中，使用者在計算圖的節點上指定函數 m:V→N m : V → N ，以訓示可以重新計算結果的次數。我們将 m m 稱之為鏡像計數函數，因為重新計算本質上是複制（鏡像）節點。當所有鏡像計數設定為0時，算法退化為正常梯度圖。在Alg. 2中指定重新計算模式時 ，使用者隻需要設定段内節點的 m(v)=1m(v)=1 ，段輸出節點的 m(v)=0 m ( v ) = 0 。鏡像計數也可以大于1，這會導緻遞歸泛化，在後續節中會進行讨論。圖3顯示了記憶體優化梯度圖的一個例子。重要的是，Alg. 2還輸出用于計算的周遊順序，是以可以優化記憶體使用。此外，這種周遊順序可以幫助依賴于運作時配置設定的架構引入控制流依賴項。

丢棄低運算成本操作的結果

通用方法的一個快速應用是丢棄低成本操作的結果并保持計算耗時的結果。這通常用于卷積神經網絡中的Conv-BatchNorm-Activation流水線。我們總是可以保留卷積結果，但是丢棄歸一化、激活和池化的結果。在實踐中，這将節省記憶體，并且計算開銷很小，因為批量歸一化和激活函數的計算成本很低。

O(n−−√) O ( n ) 記憶體代價算法

Alg. 2提供了一種計算換記憶體的通用方法。但它仍然需要了解應該保留哪個中間結果，以及要重新計算哪個結果。假設我們将 n n 層網絡劃分為 kk 段，訓練此網絡的記憶體成本如下給出。

cost-total=maxi=1,…,kcost-of-segment(i)+O(k)=O(nk)+O(k) cost-total = max i = 1 , … , k cost-of-segment ( i ) + O ( k ) = O ( n k ) + O ( k )

等式的第一部分是在每個段上反向傳播的記憶體開銷。鑒于段是平等配置設定的，這将轉化為 O(n/k) O ( n / k ) 的成本。等式的第二部分是存儲分段之間的中間輸出的成本。設定 k=n−−√ k = n ，我們得到 O(2n−−√) O ( 2 n ) 的成本。該算法隻需在訓練過程中增加一次前向，但将記憶體開銷降低為次線性。由于後向操作幾乎是前向的兩倍，是以隻會稍許減慢計算速度。

圖3 記憶體優化梯度圖生成示例。鏡像正向路徑以表示梯度計算時發生的重計算。使用者指定鏡像因子以控制是否應放棄或保留結果。

一般在多數情況下，每層的記憶體成本是不一樣的，是以我們不能簡單地設定 k=n−−√ k = n 。但是，中間輸出和每個階段成本之間的權衡仍然存在。在這種情況下，給定每個段内的記憶體開銷預算作為單一參數 B B ，我們使用Alg. 3以貪婪算法進行配置設定。不同的 B B 給出我們不同的配置設定計劃，要麼給中間輸出配置設定更多的記憶體，要麼給每個階段配置設定更多的計算。當我們進行靜态記憶體配置設定時，給定每個配置設定計劃，我們可以得到精确的記憶體消耗。我們可以使用這些資訊對 B B 進行啟發式搜尋，以找到平衡兩者成本的最佳記憶體方案。搜尋步驟的細節在補充材料中介紹。我們發現這種方法在實踐中運作良好。我們還可以推廣該算法，通過考慮執行每個操作的成本以盡可能地保留耗時的操作。

更一般化的視角：遞歸和子程式

在本節中，我們提供了上述記憶體優化方案的另一種了解。具體來說，我們可以将每個段視為一個塊操作符，它将段中的所有操作組合在一起。這一思路如圖4所示。組合運算符在描述其内部計算的子圖上運作來計算梯度。這一觀點允許我們将一系列操作視為子程式。子圖中的優化不會影響外部世界。是以，我們可以遞歸地将我們的記憶體優化方案應用于每個子圖。

圖4 記憶體配置設定優化的遞歸視圖。分段可以看作是一個單一的運算符，它将段内的所有運算符組合在一起。在每個運算符中，執行一個子圖來計算梯度。

用遞歸進一步減少記憶體 設 g(n) g ( n ) 為在 n n 層神經網絡上進行前向和後向傳播的記憶體代價。假設我們将 kk 個中間結果存儲在圖中，并在子路徑上進行前向和後向傳播時遞歸地應用相同的政策。我們有以下遞歸公式。

g(n)=k+g(n/(k+1)) g ( n ) = k + g ( n / ( k + 1 ) )

求解這個遞歸公式給了我們

g(n)=klogk+1(n) g ( n ) = k log k + 1 ⁡ ( n )

作為一個特例，如果我們設 k=1 k = 1 ，我們得到 g(n)=log2n g ( n ) = log 2 ⁡ n 。這是一個有趣的結論，因為訓練一個 n n 層神經網絡，現有的所有實作都僅需 O(n)O(n) 的記憶體用于特征映射。這将需要 O(log2n) O ( log 2 ⁡ n ) 的正向成本，因而不太可能普遍使用。但它示範了我們如何通過遞歸換取記憶體。

深度學習架構指南

在本節中，我們已經證明能夠以計算換記憶體并将其與上一節中提出的系統優化結合起來。這有助于深度學習架構

• 啟用選項以删除低成本操作的結果。
• 提供規劃算法以提供有效的記憶體方案。
• 使使用者能夠在計算圖中設定鏡像屬性以進行記憶體優化。

雖然最後一個選項不是必需的，但提供這樣的接口使得使用者可以創造他們自己的記憶體優化器，并鼓勵将來對相關方向進行研究。在這一願景下，我們支援定制圖鏡像方案，并将公開源代碼。

B預算搜尋

Alg. 3使我們可以在給定單個參數 B B 的情況下生成優化的記憶體方案。該算法依靠近似記憶體估計獲得更快的速度。在方案完成後，我們可以使用靜态配置設定算法計算确切的記憶體成本。然後，我們可以在 B B 上進行網格搜尋以找到一個好的記憶體方案。

為了獲得網格的設定，我們首先以 B=0 B = 0 運作配置設定算法，然後改用 B=xy−−√ B = x y 運作一次。這裡， x x 和 y y 是 Alg. 3 在第一次運作中的輸出。 x x 是存儲段間特征映射的近似成本， y y 是運作每段的記憶體近似成本。 B=xy−−√ B = x y 為每一階段記憶體成本的估計。至此已經可以提供一個很好的記憶體方案。然後，我們圍繞 B=xy−−√ B = x y 設定網格，以進一步優化解決方案。

在實踐中，我們發現在 [B/2–√,2–√B] [ B / 2 , 2 B ] 上使用大小為 6 6 的網格可以給出實驗中良好的記憶體方案。我們用python實作了配置設定算法，而未嘗試優化速度。我們的代碼僅需幾秒鐘就能獲得實驗所需的方案。

實驗

實驗結果及完整内容請參考原文

算法實作

下面是福利時間。對于MXNet使用者，作者已經公布代碼；而TFBOYS可以使用OpenAI的gradient-checkpointing。

影響力

論文的價值可以通過引用量間接展現。例如引證文獻：

MegDet: A Large Mini-Batch Object Detector

Memory-Efficient Implementation of DenseNets