在過去的十年中，深度學習的一個持續的趨勢是向更深更大的網路發展。儘管硬體效能迅速提高，但先進的深度學習模型仍在不斷挑戰GPU RAM的極限。因此，即使在今天，人們仍然希望找到一種方法來訓練更大的模型，同時消耗更少的記憶體。這樣做可以讓我們更快地進行訓練，使用更大的批處理大小，從而實現更高的GPU利用率。

在本文中，我們探討了為深度神經網路優化記憶體分配的技術。我們討論了幾個備選解決方案。雖然我們的建議並非詳盡無遺，但這些解決方案是有益的，可以讓我們介紹主要的設計問題。

1 計算圖

首先，讓我們回顧一下計算圖的概念。計算圖描述了深度網路中操作之間的(資料流)依賴關係。圖中執行的操作可以是細粒度的，也可以是粗粒度的。下圖顯示了兩個計算圖示例。

計算圖的概念被顯式地編碼在像Theano和CGT這樣的包中。在其他庫中，計算圖隱式地顯示為網路配置檔案。這些庫的主要區別在於它們如何計算梯度。主要有兩種方法:在同一圖上執行反向傳播，或者顯式地表示反向路徑來計算所需的梯度。

Caffe、CXXNet和Torch等庫採用了前一種方法，在原始圖形上執行支援。像Theano和CGT這樣的庫採用後一種方法，顯式地表示向後路徑。在這個討論中，我們採用顯式的反向路徑方法，因為它具有一些易於優化的優點。

但是，我們應該強調的是，選擇顯式的向後路徑方法並不侷限於符號庫，比如Theano和CGT。我們還可以使用顯式的反向路徑來計算基於層的庫(它將前向和後向連線在一起)的梯度。下圖展示瞭如何做到這一點。基本上，我們引入一個反向節點，它連結到圖的正向節點並呼叫該層。在反向操作中呼叫layer.backward。

這個討論幾乎適用於所有現有的深度學習庫。(庫之間存在差異，例如高階微分，但這超出了本主題的範圍。)

為什麼顯式向後路徑更好?讓我們用兩個例子來解釋它。第一個原因是顯式的向後路徑清楚地描述了計算之間的依賴關係。考慮下面的情況，我們想要得到A和b的梯度，從圖中我們可以清楚地看到，d©梯度的計算並不依賴於F，這意味著我們可以在完成正向計算後立即釋放F的記憶體。同樣，C的記憶體也可以回收。

顯式反向路徑的另一個優點是能夠擁有不同的反向路徑，而不是正向路徑的映象。一個常見的例子是分割連線，如下圖所示。

在本例中，B的輸出由兩個操作引用。如果我們想在同一個網路中做梯度計算，我們需要引入一個顯式的分割層。這意味著我們也需要對前向計算進行分割。在這個圖中,前進計算並不包含一個分離層,但圖會在梯度回傳到B之前自動插入一個梯度聚合節點,這幫助我們節省前向計算中的輸出分離層的記憶體佔用和資料複製開銷。

如果我們採用顯式的向後傳遞方法，向前傳遞和向後傳遞之間沒有區別。我們只是按照時間順序一步一步地進行計算。這使得顯式的向後方法易於分析。我們只需要回答這個問題:如何為計算圖的每個輸出節點分配記憶體？

2 什麼是可以優化的

如您所見，計算圖是討論記憶體分配優化技術的一種有用方法。我們已經展示瞭如何使用顯式的後向圖來節省一些記憶體。現在，讓我們進一步研究優化，並看看如何為基準測試確定合理的基線。

假設我們要建立一個有n層的神經網路。通常，在實現神經網路時，我們需要為每一層的輸出和反向傳播過程中使用的梯度值分配節點空間。這意味著我們需要大約2n個記憶體單元。在使用顯式向後圖方法時，我們面臨相同的要求，因為向後傳遞中的節點數量與向前傳遞中的節點數量大致相同。

2.1 就地操作

我們可以使用的最簡單的技術之一是跨操作的就地記憶體共享。對於神經網路，我們通常可以將該技術應用於與啟用函式相對應的操作。考慮以下情況，我們希望計算三個鏈式sigmoid函式的值。

因為我們可以就地計算sigmoid，使用同一塊的記憶體儲存輸入和輸出，我們可以使用常量記憶體計算任意長度的sigmoid函式鏈。

注意:在實現就地優化時很容易出錯。考慮以下情況，其中B的值不僅由C使用，而且由F使用。

我們無法進行就地優化，因為在計算C=sigmoid(B)之後仍然需要B的值。對每個sigmoid操作進行就地優化的演算法可能會落入這樣的陷阱，所以我們需要小心何時可以使用它。

2.2 標準的記憶體共享

就地操作並不是我們可以共享記憶體的唯一地方。在下面的例子中，由於計算E後不再需要B的值，我們可以重用B的記憶體來儲存E的結果。

記憶體共享不一定需要相同的資料形狀。注意，在前面的例子中，B和E的形狀可能不同。為了處理這種情況，我們可以分配一個大小等於B和E所需要的最大記憶體區域，並在它們之間共享它。

2.3 實際神經網路分配的例子

當然，這些只是玩具的例子，它們只涉及前向計算。但同樣的想法也適用於真實的神經網路。下圖顯示了一個兩層感知器的分配計劃。
在本例中：

• act1、d(fc1)、out、d(fc2)計算採用就地優化；
• d(act1)和d(A)之間共享記憶體。

3 記憶體分配演算法

到目前為止，我們已經討論了優化記憶體分配的一般技術。我們已經看到了需要避免的陷阱，在就地記憶體優化的例子中已經證明了這一點。那麼，我們如何正確分配記憶體呢?這不是一個新問題。例如，它非常類似於編譯器中的暫存器分配問題。也許我們可以借鑑一些技術。我們在這裡並不是要對技術進行全面的回顧，而是要介紹一些簡單但有用的技巧來解決這個問題。

關鍵問題是我們需要放置資源，這樣它們才不會相互衝突。更具體地說，每個變數從計算到最後一次使用之間都有一個生命週期。在多層感知器的情況下，計算act1後fc1的生命週期結束。

其原則是隻允許在生命週期不重疊的變數之間共享記憶體。有多種方法可以做到這一點。您可以構造每個變數作為節點的衝突圖，並將具有重疊生命週期的變數之間的邊連結起來，然後執行圖形著色演算法。這可能具有 $O(n^2)$複雜度，其中 $n$是圖中節點的數量。這種方法計算量太大。

讓我們考慮另一個簡單的啟發式。其思想是模擬遍歷圖的過程，並對依賴於節點的未來操作進行計數。

• 只有當前操作依賴於源(即count等於1)的時候，才可以使用就地優化；
• 只有當計數變為0時，記憶體可以回收到右上角的框中；
• 當我們需要新的記憶體時，我們可以從右上角的框中獲取一個已釋放的記憶體，或者分配一個新的記憶體。
  注意:在模擬過程中，沒有分配記憶體。相反，我們記錄每個節點需要多少記憶體，並在最終的記憶體計劃中分配最大的共享部分。

3.1 靜態分配 VS. 動態分配

前面的策略完全模擬了命令式語言(如Python)中的動態記憶體分配過程。計數是每個記憶體物件的引用計數器，當引用計數器變為0時，物件將被垃圾收集。從這個意義上說，我們模擬動態記憶體分配一次，以建立一個靜態分配計劃。我們可以簡單地使用動態分配和釋放記憶體的命令式語言嗎?

主要區別在於靜態分配只完成一次，因此我們可以使用更復雜的演算法。例如，我們可以搜尋與所需記憶體塊相似的記憶體大小。還可以使分配具有圖意識。我們將在下一節討論這個問題。動態分配給快速記憶體分配和垃圾回收帶來了更大的壓力。

對於希望依賴動態記憶體分配的使用者來說，還有一點需要注意:不要不必要地引用物件。例如，如果我們組織一個列表中的所有節點，然後將它們儲存在一個Net物件中，這些節點將永遠不會被解除引用，我們也不會獲得任何空間。不幸的是，這是一種常見的組織程式碼的方法。

3.2 為並行操作分配記憶體

在上一節中，我們討論瞭如何模擬執行計算圖的過程來獲得靜態分配計劃。然而，優化平行計算也帶來了其他挑戰，因為資源共享和並行一般不可兼顧。讓我們看一下同一個圖的下面兩個分配計劃。
如果我們進行從[1]到[8]的序列計算，這兩個分配計劃都是有效的。然而，左邊的分配計劃引入了額外的依賴項，這意味著我們不能並行執行[2]和[5]的計算。右邊的計劃可以。為了使計算並行化，我們需要更加小心。

3.3 正確和安全是第一位的

正確是我們的首要原則。這意味著以一種考慮到隱式依賴記憶體共享的方式執行。您可以通過向執行圖中新增隱式依賴項邊來實現這一點。或者，更簡單的是，如果執行引擎能夠識別可變變數，就像我們在依賴引擎設計的討論中所描述的那樣，按順序推進操作，並寫入表示相同記憶體區域的相同變數標記。

首先總是要使用一個安全的記憶體分配計劃。這意味著永遠不要將相同的記憶體分配給可以並行化的節點。當更需要減少記憶體時，這可能並不理想，而且當我們可以從同時在同一個GPU上執行的多個計算流中獲益時，也不會得到太多好處。

3.4 嘗試允許更多的並行化

現在我們可以安全地執行一些優化。一般的想法是嘗試並鼓勵不能並行化的節點之間的記憶體共享。您可以通過建立一個依賴關係圖並在分配期間查詢它來實現這一點，這將在構建時花費大約 $O(n^2)$的時間。我們還可以在這裡使用啟發式，例如，為圖中的路徑著色。如下圖所示，當您試圖找到圖中最長的路徑時，將它們塗上相同的顏色並繼續。
在獲得節點的顏色之後，只允許(或鼓勵)相同顏色的節點之間共享。這是祖先關係的更嚴格版本，但是如果只搜尋第一個k路徑，只需要O(n)時間。

4 我們能節省多少記憶體呢？

我們已經討論了一些技術和演算法，您可以使用它們來壓縮深度學習的記憶體。但使用這些技術可以真正節省多少記憶體呢?

對於已經為大型操作優化過的粗粒度操作圖，可以將記憶體消耗減少大約一半。如果您正在優化符號庫(如Theano)使用的細粒度計算網路，則可以進一步減少記憶體使用。本文的大部分思想啟發了MXNet的設計。

此外，您還會注意到，與同時執行前向和後向傳遞相比，僅執行前向傳遞的記憶體成本非常低。這是因為如果只執行前向傳遞，記憶體重用會更多。

這裡有兩個結論:

• 使用計算圖來分配記憶體；
• 對於深度學習模型，預測比訓練消耗更少的記憶體。

文章翻譯自：https://mxnet.incubator.apache.org/architecture/note_memory.html