原文地址:https://github.com/dmlc/mxnet/blob/master/docs/zh/note_memory.md

最近在學習mxnet，發現mxnet的有些思想挺好，這裡試著翻譯它的一篇文章，也試著瞭解減少記憶體消耗的設計原則，希望這篇文章對大家有幫助。原文連結Squeeze the Memory Consumption of DeepLearning

深度學習的重要主題是關於訓練更深度和更大型的網路。最近幾年，硬體普遍升級的相當迅速，這種巨型的深度網路怪物常常對視訊記憶體有更多的需求。如果同樣的網路模型我們能使用更少的記憶體意味著我們每批輸入資料可以輸入更多，也能增加GPU的利用率。

這篇文章討論了深度神經網路中如何對記憶體分配優化，並且提供了一些候選的解決方案。文章裡所討論的解決方案並不意味著全部完成了，但我們想其中的例子能對大部分情況有用。

計算圖譜（圖計算）

我們將從計算圖譜開始討論，這些篇幅的後面部分有一些工具可以幫助我們。圖計算描述（資料流）在深度網路中的運算依賴關係。圖裡面的運算執行可以是細粒度也可以是粗粒度。以下給出計算圖譜的兩個例子。

Comp Graph Example

計算圖譜的這些概念在如Theano，CGT的包中有深入體現。事實上，它們也隱含的存在於大多數網路結構的庫中。這些庫的主要區別是他們如何做梯度計算。這裡有兩種方法，在同一個圖裡用BP演算法，或者用一個顯式反向路徑（explicitbackward path）計算所需梯度。

Backward Graph

像caffe, cxxnet, torch這些庫在同一個圖中使用了BP演算法。而像Theano, CGT這些庫使用了顯式反向路徑逼近演算法。這篇文章中我們採用了顯式反向路徑，因為該方法在輪流的優化中可以帶來更多的優勢。

然而，我們要強調的是，選擇顯式反向路徑法的執行並不限制我們使用Theano,CGT庫的範圍。我們也可以將顯式反向路徑法用在基於層（前向，後向在一起）的庫做梯度計算。下面的圖表展示如何做到。基本上，我們可以引出一個反向節點到圖的正向節點，並且在反向的操作中可以呼叫layer.backward。

Backward Layer

這次討論中提供幾乎所有已存在的深度學習庫（這些庫有些區別，比如高階分化，這些不在本篇文章中討論）

為什麼明確反向路徑更好?讓我們用兩個例子來解釋吧。第一個原因是,明確反向路徑能清晰描述計算間的依賴關係。考慮下面的情況,我們想要得到的A和B的梯度。我們可以從圖中清楚地看到,d(C)的梯度計算不依賴於F,這意味著我們可以在正向計算完成後釋放記憶體,同樣C的記憶體也可以回收。

Backward Prune

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

Backward Agg

在這個例子中,兩個操作引用了B的輸出。如果我們想在同一網路做梯度計算,一個顯式分割層需要引入。這意味著我們需要分割正向傳遞。在這個圖中,正向傳遞不包含一個分割層，但圖會在通過梯度回到B之前，自動插入一個梯度聚合節點。這幫助我們節省輸出分割層的記憶體的分配，以及在正向傳遞中操作複製資料的消耗。

如果我們採用計算圖的顯式反向檢視,在正向和反向傳遞中是沒有區別的。我們只會簡單的按照計算圖的拓撲順序進行,並完成計算。這也簡化了討論。現在的問題就變成:

如何在每一個計算圖的輸出節點上分配記憶體?

嗯,似乎與深度學習無關,但更多的上下文編譯、資料流優化等有關。但它確實是深度學習的飢餓怪獸,激勵我們解決這個問題,並從中受益。

如何優化？

希望你能相信，圖計算是記憶體分配優化技術中的好方法。正如你所看到的顯式反向圖可以已經可以節省一些記憶體。讓我們討論下，我們可以做什麼優化和底線是什麼。

假設我們要建設具有n層的神經網路。神經網路的一個典型實現是需要給每一層的輸出分配節點的空間，以及反向傳播梯度值。這意味著我們需要大致2n個記憶體單元。這在顯式反向圖中也一樣，在反向傳播節點與前向傳播節點大致相同。

替代操作

其中第一個我們可以做的事就是替代操作記憶體共享。這通常是簡單的操作，如啟用函式。考慮以下情況，在這裡我們要計算3個chainedsigmoid函式的值。

Inplace op

因為我們可以用代替方式計算sigmoid函式，即，為輸入輸出使用相同的記憶體區。我們可以簡單地分配一份記憶體，並用它計算sigmoid鏈的任意長度。

然而，替代優化有時也是一種錯誤的方式，特別是這個包變的有點普通。考慮以下情況，其中，B的值不是僅由C也由F使用。

Inplace trap

因為B的值在C=sigmoid(B)後仍然需要，我們不能執行替代優化。所以為每一個sigmoid操作做簡單替代優化演算法可能會落入這個陷阱，我們在做的時候必須要小心。

普通的記憶體共享

除了替代操作，記憶體也可以共享。考慮以下情況下，由於在我們計算了E之後，B的值是不再需要，我們可以重新使用記憶體來儲存E的結果。

Normal Sharing

我們想指出的是在相同的資料模型上記憶體共享並不是必須的。在上面的例子中，B和E的模型可以是不同的，我們可以簡單地按兩個最大的尺寸分配記憶體區域，也可以在兩者之間共享它。

實時神經網路分配示例

上面的例子都是構成情況，僅包含正向傳遞的計算。實際上這個思想也可以用於儲存實時的神經網路。下圖顯示了分配方案：我們可以為一個兩層感知器分配。

Net Alloc

在上面的例子：

替代優化會被應用在計算act1，d(fc1)，out和d(fc2)。記憶體共享被用在d(act1)和d(A)之間。記憶體分配演算法

上一節中我們已經討論瞭如何用一般的技術來優化記憶體分配。但是，我們也看到一些陷阱，類似於替代這是我們想要避免的情況。我們怎樣才能正確地分配記憶體？這不是一個新的問題。例如，它與編譯器暫存器分配非常相似。所以，我們可以借鑑很多思想。我們不打算在這裡給出全面的技術回顧，而是介紹一些簡單實用的技巧來解決問題。

關鍵問題是，我們要放置資源，使得它們不互相沖突。更具體地講，每個變數有一個life time。它在計算時間和上個使用時間之間。在多層感知器的例子中，fc1的life time在act1計算後結束。

Net Alloc

其原理是，僅允許變數在不重疊的生命週期中共享記憶體。有多種方式來解決這個問題。一種可能的方法是構造一個衝突的圖，這個圖模型每一個變數作為節點和連結邊在一些具有重疊生命週期變數中，並執行一個圖著色演算法。這可能需要$O(n^2)$的時間複雜度，其中n是圖中的節點的數量，這也是較為合理的消耗。

在這裡我們將介紹另一種簡單的有啟發的方法。這樣做是為了模擬遍歷圖的步驟，並保持未來依賴於節點的操作計數器。

Alloc

只有當前的操作依賴於來源時，就可以執行一個替代優化（即，counter=1）當計數器變為0時，記憶體可以回收到在右上角的框中每一次，當我們需要新的記憶體，我們可以從框中得到它，或分配一個新的。

其中值得注意的是，在模擬過程中，將不會分配記憶體，但我們可以保留每個節點需要多少記憶體的資訊，並在最終的記憶體方案中分配最大共享部分。

靜態VS動態分配

如果你仔細想想，在如python的指令碼語言中，你會發現上面的策略正好模擬了動態記憶體分配的過程。這個計數器是每一個記憶體物件的引用計數器，當引用計數器變為零物件被回收。在這個意義上，我們模擬了動態記憶體分配一次，就建立一個靜態分配方案。現在的問題是，我們可以簡單地使用指令碼語言去動態分配並釋放記憶體嗎？

主要的區別是靜態分配只做一次，因此，我們可以使用更復雜的演算法

例如，在記憶體中搜索的情形和請求記憶體塊是比較相似的。圖模型也可以意識到這種分配，見下一節更多的討論。動態方式將會給快速記憶體分配器和垃圾收集器帶來更大的壓力。

此外，對那些想要回應動態記憶體分配的使用者來說，也有另一種方法：不要獲取不必要的物件引用。例如，在一個網路中，如果我們組織和儲存一個列表中所有的節點，這些節點將一直有個引用（reference），讓我們沒有可以利用的空間分配記憶體。不幸的是，這是組織程式碼的一種常見方法。

在並行操作上分配

在上一節中，我們討論我們如何模擬圖計算的執行步驟，並得到一個靜態分配方案。不過，在我們使用並行優化計算時有更多的問題，資源共享和並行是天平的兩端。讓我們來看看在同一圖形下的兩個分配方案：

Parallel Alloc

如果我們以序列方式執行從A[1]到A[8]，那這兩種分配方案是都有效的。然而，左側的分配方案引入更多的依賴，意味著我們不能以並行方式執行A[2]和A[5]的計算，而右邊可以執行。

我們可以看到，如果我們想要平行計算，需要更多的關注計算方面的關係。

優先保證安全和正確

保正正確，這是我們需要知道的第一原則。這意味著執行的這種方式需要考慮採取隱式依賴記憶體共享。這可以通過新增隱式依賴邊到執行圖模型來完成。或者更簡單的，按照依賴引擎說明dependency engine note中所述，如果執行引擎意識到變化了，往序列裡推送一個操作和寫入同樣的變數標記，這個標記表示相同的記憶體區域。

另一種方式是一直產生記憶體分配方案，這是安全的，也意味著永遠都不分配同樣的記憶體給可以並行的節點。可能這不是理想的方式，因為有時記憶體的減少是更可取的，我們可以在相同的GPU獲取由多個計算流的執行結果，但這沒有太大的提高。

儘量允許更多的並行

假定我們是正確的，我們現在安全的做了一些優化。總的想法是儘量鼓勵在無法並行節點之間做記憶體共享。這又可以在分配時通過建立和查詢一個父輩關係圖，它耗費約$O(n^2)$時間來構建完成。這裡我們也可以使用探索法，例如，可以將圖上的路徑標記好顏色。這個想法在下圖所示，每次我們嘗試找到一個圖中最長路徑，給他們標記相同的顏色，然後繼續。

Path Color

我們得到的節點的顏色之後，可以在同樣顏色的節點之間共享（或鼓勵這種共享）記憶體。這是一個比父輩關係更精確的版本，如果我們只搜尋第一個k路徑它的時間成本為$O(n)$。

這裡討論的策略是絕不是唯一的解決方案，我們希望能沿著這條路找到更多的高階方法。

我們可以節省多少記憶體？

感謝您閱讀到這部分！我們已經討論了可以用來壓縮深度學習的記憶體使用量的技術和演算法。現在問題來了，我們可以用這些技術節省多少記憶體？

答案是，我們可以通過使用這些技術大致減少記憶體消耗一半。這是已經優化的粗粒度操作圖大操作。如果我們要優化一個更細粒度的計算網路，可以使用如Theano這樣的庫，將減少更多的記憶體消耗。

本文的大部分想法激發mxnet的設計。我們提供了一個記憶體消耗估算指令碼Memory Cost Estimation Script，你可以試一下，看看我們在不同的策略下需要多少記憶體。

如果你要試試這個指令碼，有一個選項叫forward_only，顯示了只執行一個正向傳遞時的消耗。你會發現這個記憶體消耗極低。如果你讀文章的前面部分，你不會感到驚訝，這是因為如果我只進行正向傳遞，有更多的記憶體被重複使用。因此，這裡有兩個結論：

使用圖計算圖巧妙和正確地分配記憶體。執行深度學習預測的消耗比深度學習訓練的記憶體消耗要少得多。貢獻說明

本說明是我們的提供給開源深度學習庫系統設計說明的一部分。你可以通過提交pullrequest為此文件做貢獻，我們非常歡迎！