本文主要译介自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
- 前向运算各层响应： x , h , y x, h, y x,h,y
这样，可以在后向运算中用梯度下降更新参数：
Δ 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)⊙Rh = 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。
之后，把 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
输入输出为多通道时，方法类似，详情参见这篇博客。
在计算此类卷积时，前层响应 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本身是一家机器学习芯片初创公司，行文中难免夹带私货，请明辨。
