【深度學習】為什么深度學習需要大內存？
本文主要譯介自Graphcore在2017年1月的這篇博客: Why is so much memory needed for deep neural networks。介紹了深度學習中內存的開銷，以及降低內存需求的幾種解決方案。 為便於閱讀，本文修改了原文分段，並添加更詳細的計算說明。
深度學習的內存消耗在哪里？
回顧：簡單例子
考慮一個單層線性網絡，附帶一個激活函數：
h
=
w
1
x
+
w
2
h=w_1x+w_2
h=w1x+w2
y
=
f
(
h
)
y=f(h)
y=f(h)
代價函數：
E
=
∣
∣
y
−
y
‾
∣
∣
2
E=||y-\overline{y}||^2
E=∣∣y−y∣∣2
在訓練時，每一個迭代要記錄以下數據：
w
1
,
w
2
w_1,w_2
w1,w2</li><li>前向運算各層響應：
x
,
h
,
y
x, h, y
x,h,y</li>
這樣，可以在后向運算中用梯度下降更新參數：
Δ
w
1
=
η
⋅
∂
E
∂
w
1
=
η
⋅
2
(
y
−
y
‾
)
⋅
f
′
(
h
)
⋅
x
\Delta w_1=\eta\cdot \frac{\partial E}{\partial w_1}=\eta \cdot 2(y-\overline{y})\cdot f'(h) \cdot x
Δw1=η⋅∂w1∂E=η⋅2(y−y)⋅f′(h)⋅x
Δ
w
2
=
η
⋅
∂
E
∂
w
1
=
η
⋅
2
(
y
−
y
‾
)
⋅
f
′
(
h
)
\Delta w_2=\eta\cdot \frac{\partial E}{\partial w_1}=\eta \cdot 2(y-\overline{y})\cdot f'(h)
Δw2=η⋅∂w1∂E=η⋅2(y−y)⋅f′(h)
內存消耗的三方面
輸入數據
很小，不做考量。
256256的彩色圖像：25625631 byte= 192KB
模型參數
較大，和模型復雜度有關。
入門級的MNIST識別網絡有6.6 million參數，使用32-bit浮點精度，占內存：6.6M 32 bit = 25MB 50層的ResNet有26 million參數，占內存：26M 32 bit = 99MB
當然，你可以設計精簡的網絡來處理很復雜的問題。
各層響應
較大，同樣和模型復雜度有關。
50層的ResNet有16 million響應，占內存：16M*32bit = 64MB
響應和模型參數的數量並沒有直接關系。卷積層可以有很大尺寸的響應，但只有很少的參數；激活層甚至可以沒有參數。
– 這樣看起來也不大啊？幾百兆而已。
– 往下看。
batch的影響
為了有效利用GPU的SIMD機制，要把數據以mini-batch的形式輸入網絡。
如果要用32 bit的浮點數填滿常見的1024 bit通路，需要32個樣本同時計算。
在使用mini-batch時，模型參數依然只保存一份，但各層響應需要按mini-batch大小翻倍。
50層的ResNet，mini-batch=32，各層相應占內存：64MB*32 = 2GB
卷積計算的影響
設
H
×
W
H\times W
H×W的輸入圖像為
X
X
X，
K
×
K
K\times K
K×K的卷積核為
R
R
R，符合我們直覺的卷積是這樣計算的。
對每一個輸出位置，計算小塊對位乘法結果之和。
Y
(
h
,
w
)
=
∑
X
k
,
k
s
(
h
,
w
)
⊙
R
Y(h,w) = \sum{X^s_{k,k}(h,w) \odot R}
Y(h,w)=∑Xk,ks(h,w)⊙R</p>
h = 1 : H , w = 1 : W h=1:H, w=1:W h=1:H,w=1:W
其中， X k , k s ( h , w ) X^s_{k,k}(h,w) Xk,ks(h,w)表示輸入圖像中，以 h , w h,w h,w為中心，尺寸為 K × K K\times K K×K的子圖像。
但是，這種零碎運算很慢。
在深度學習庫中，一般會采用lowering的方式，把卷積計算轉換成矩陣乘法。
首先，把輸入圖像分別平移不同距離，得到
K
2
K^2
K2個
H
×
W
H\times W
H×W的位移圖像，串接成
H
×
W
×
K
2
H\times W \times K^2
H×W×K2的矩陣
X
‾
\overline{X}
X。<br> 之后，把
K
×
K
K\times K
K×K的卷積核按照同樣順序拉伸成
K
2
×
1
K^2\times 1
K2×1的矩陣
R
‾
\overline{R}
R
Y
=
X
‾
⋅
R
‾
Y=\overline{X}\cdot \overline{R}
Y=X⋅R</p>
輸入輸出為多通道時，方法類似，詳情參見這篇博客。
在計算此類卷積時，前層響應
X
X
X需要擴大
K
2
K^2
K2倍。
50層的ResNet，考慮lowering效應時，各層響應占內存7.5GB
使用低精度不能降內存
為了有效利用SIMD，如果精度降低一倍，batch大小要擴大一倍。不能降低內存消耗。
降內存的有效方法
in-place運算
不開辟新內存，直接重寫原有響應。 復雜一些，通過分析整個網絡圖，可以找出只需要用一次的響應，它可以和后續響應共享內存。例如MxNet的memory sharing機制。
綜合運用這種方法，MIT在2016年的這篇論文能夠把內存降低兩到三倍。
計算換存儲
找出那些容易計算的響應結果（例如激活函數層的輸出）不與存儲，在需要使用的時候臨時計算。
使用這種方法，MxNet的這個例子能夠把50層的ResNet網絡占用的內存減小四倍。
類似地，DeepMind在2016年的這篇論文用RNN處理長度為1000的序列，內存占用降低20倍，計算量增加30%。
百度語音在2016年的這篇論文同樣針對RNN，內存占用降低16倍，可以訓練100層網絡。
當然，還有Graphcore自家的IPU，也通過存儲和計算的平衡來節約資源。
Graphcore本身是一家機器學習芯片初創公司，行文中難免夾帶私貨，請明辨。