Deep Nets Sublinear Memory Cost

论文链接：arXiv:1604.06174

Abstract

设计一个算法，训练一个N层网络仅耗费 O ( n ) O(\sqrt{n}) O(n )的内存，每个mini-batch只需要一个额外的前向计算成本，因为许多先进的模型已经达到了GPU显存的上限，需要去探索更深入更复杂的模型，推进深度学习研究的创新。

这个算法专注于降低训练期间存储中间特征图和梯度的内存成本，计算图（DAG）分析用于automatic in-place operation和内存共享优化，这个训练算法对内存利用更高效只需要一点额外的计算成本，可以通过计算来换取内存资源。

极端情况下，内存消耗可以减少到 O ( l o g n ) O(log\ n) O(log n)，而正向计算的额外开销只需要 O ( n l o g n ) O(n\ log\ n) O(n log n)，这种算法能将1000层深度残差网络的内存成本从48G降低到7G，类似地，在非常长的序列中训练复杂递归神经网络时也能显著的降低内存成本。

总的来说，就是用计算时间换空间

1 Introduction

这类系统性方法用于减少深度神经网络训练时的内存消耗，其主要关注于减少存储中间结果（特征映射）和梯度的内存成本。在许多常见的深度架构中，与中间特征映射的大小相比，参数的大小相对较小。

我们使用计算图分析来执行automatic in-place operation和内存共享优化，之后我们用一个更实用的算法，将花费O(log n)的内存用于特征映射，这个仅仅用双倍正向计算成本来训练n层网络，更为突出的，在极端情况下，可以使用O(log n)的内存存储特征映射来训练一个n层网络。

使用这类算法的原因是，目前的很多模型的趋势是用更深的架构来捕捉大量训练数据中的复杂模式，由于存储特征映射及其梯度的成本随网络深度线性变化，因此我们探索深层模型的能力受到设备（通常为GPU）内存的限制。

如当前最先进的一个模型已经耗尽了内存，从长远来看，理想的ML系统应该能够从越来越多的训练数据中不断学习，因为最佳模型大小和复杂性通常随着训练数据增多而增加，因此节约内存的训练算法是非常重要的。

减少内存消耗不仅使我们能够训练更大的模型，它还可以实现更大的批量大小，从而提高批量操作符的设备利用率和稳定性，如批量归一化。对于内存有限的设备，它有助于改善内存局部性，并可能产生更好的内存访问模式，它还使我们训练DNN时能够从模型并行性切换到数据并行性，这在某些情况下可能是有益的。

这篇文章的解决方案使学习者能够训练更深的卷积神经网络，以及具有更大unrolling steps（时序展开）的RNN，根据本文提出的内存优化技术，可以为深度学习框架提供指导。

内存优化算法实现：MXNET Memory Monger

2 Memory Optimization with Computation Graph

回顾一下计算图的概念和内存优化技术，其中一些技术已经被现有的框架如Theano，Tensorflow和MXNet所使用，计算图由表示操作之间依赖关系的nodes和edges组成。

下图为一个two-layer fully connected NN计算图的例子，这是使用粗粒度的向前和向后操作来使图更简单，我们通过隐藏权重节点和权重的梯度来进一步简化图，实践中使用的计算图可能更复杂，并且混合包含粗/细粒度操作，下图的分析可以直接用于那些更一般的情况。

上图为两层FNN训练过程的计算图和可能的内存分配计划

每个节点代表一个操作，每个边表示操作之间的依赖关系，具有相同颜色的节点共享内存以存储输出或反向传播的梯度。

为使图更清楚，这里省略了图中的权重以及输出梯度的节点，并假定了在后向运算时会计算权重的梯度。

一旦给出了NetWork Configurations（forward graph），我们就能够构建相应的backward graph进行梯度计算，backward pathway由逆向拓扑顺序（reverse topological order）配置构建，使得训练中的梯度计算步骤简化为了一个正向传播。

在训练DCN/DRN时，通常会使用大部分内存来存储中间输出和梯度，这些中间结果中的每一个都对应了图中的一个节点。智能分配算法能尽可能通过内存共享来为这些节点分配最少量的内存，如图1所示，我们可以使用两种类型的内存优化：

in-place operation： 直接将输出值存储到输入值的内存中
memeory-sharing： 中间结果不再需要的内存可以循环利用并提供给其他的节点

再看Fig.1的分配计划，第一个Sigmoid转换通过in-place来节省内存，而其backward operation又进行了重用；softmax gradient又能与第一个fullc-backward共享内存。

但是在特定情况下应用这些优化也会导致错误，比如一个operation的操作仍然是另一个操作所需要的，那么对这个输入进行in-place会导致错误的结果。

因此我们只能在生命周期lifetime没有overlap的节点之间共享内存，有多种方法可以解决这个问题：

选择以每个变量作为节点，变量之间的重叠生命周期为边，构造conflicting graph，然后对图进行graph-coloring，这将花费 O ( n 2 ) O(n^2) O(n2)的计算时间。这里采用更简单的启发式方法（simple heuristic），仅需 O ( n ) O(n) O(n)的时间。

图2以拓扑顺序遍历图，并使用计数器来指示每条记录的活性，当没有其他操作等待其输入时，可能会产生in-place操作。另一个节点使用recycled tag进行内存共享，这也可以作为图遍历的dynamic runtime algo.，并使用垃圾回收器（garbage collector）来回收过时的内存。

我们将其用作静态内存分配算法，在执行开始之前为每个node分配内存，以避免在运行时产生垃圾收集的开销。

关于图2，计算图上的内存分配算法，每个节点都关联一个liveness counter，用于记录要填充的操作。temporal tag用于指示内存共享，如果当前操作是剩下的唯一操作（输入的计数器等于1时），则可以执行原地操作，当node‘s counter记为0时，可以回收节点的标签。

问题：数据依赖性会导致每个输出的使用期更长，并且增加了大型网络的内存消耗

对这个DL框架来说非常重要的是：

以minimum manner 声明 gradient operator 的依赖性要求
对 dependency info 应用 liveness analysis 并启用内存共享

声明最小的依赖关系很重要，如果sigmoid-backward也取决于第一个fullc-forward的输出，那么fig1的allocation是不可能实现的。

Dependence analysis通常可以将n层网络预测的内存占用从 O ( n ) O(n) O(n)减少到接近 O ( 1 ) O(1) O(1)，因为每个中间结果之间可以进行共享，这样有助于减少训练的内存占用，尽管只是一个常数因素。

3 Trade Computation for Memory

3.1 General Method

虽然2中的技术可以减少DNN的训练和预测的内存占用，但是因为大多数的梯度算子都依赖于正向通路的中间结果，因此训练n层卷积网络或序列长度为n的RNN，仍然需要O(n)的内存用于中间结果。

因此为了进一步减小内存，需要去删除一些中间结果，并在需要的时候借助额外的前向计算恢复他们，更具体地说，在反向传播阶段，我们可以通过从最近的记录结果运行前向计算，来获得丢弃的中间结果。

下面是一个直链前馈NN的简化算法，具体来说，NN被分为几个部分，这个算法只记住每段的输出并删除段内的所有中间结果。反向传播期间在段内重新计算丢弃的结果，因此我们的内存开销只需要覆盖所有段的输出以及逐段反向传播时的最大内存占用。

只要将图分成几部分，Alg.1也可以推广到其他常用的计算图。但是直接应用Alg.1有两个缺点：

用户必须手动划分图写定制的train loop
将无法使用上节中介绍的其他内存优化方法

Alg.2通过引入General gradient graph construction alg.来解决这个问题。

Alg.2中，用户在computation graph的节点上指定函数m： V → N V\rightarrow N V→N，这个指示可以重新计算结果的次数。将m记为mirror count funtion，因为recomputation本质上是duplicate mirror node，当所有mirror计数设置为0时，算法将退化为normal gradient graph。

在Alg.2中指定指定recomputation模式时，用户只需要设置段内节点的 m ( v ) = 1 m(v)=1 m(v)=1，段输出节点的 m ( v ) = 0 m(v)=0 m(v)=0，mirroc count也可以大于1，但这会导致递归泛化。

Fig3 内存优化梯度图，the forward path is mirrored to represent the re-computation happened at gradient calculation。用户可以置顶镜像因子来控制是否应该放弃或者保留结果

多数情况下，每层的内存成本是不一样的，所以我们不能简单的设置 k = n k=\sqrt{n} k=n 。但是，中间输出和每个阶段成本之间的权衡仍然存在，在这种情况下，给定每个段内的内存开销预算作为单一参数B。

引入Alg.3以贪心进行分配，不同的B给出不同的分配计划，要么给中间输出分配更多的内存，要么给每个阶段分配更多的计算。当我们进行静态内存分配的时候，给定每个分配计划，我们可以得到精确的内存消耗，我们可以使用这些信息对B进行heuristic search，来找到平衡两者成本的最佳内存方案。

3.2 Recursion & Subroutine

在这里提供了内存优化方案的另一种理解，具体来说，可以将每个段视为一个块操作符，将段中所有操作组合在一起。组合运算符在描述其内部计算的子图上来计算梯度。这一观点允许我们将一系列操作视为子图，子图中的优化不会影响外部世界，因此我们可以递归地将我们的内存优化方案应用于每个子图。

Fig4 内存分配优化的递归视图，分段可以看作是一个单一的运算符，它将段内的所有运算符组合在一起，每个运算符中，执行一个子图来计算梯度。

**Pay even less memory with recursion **

设 g ( n ) g(n) g(n)为在n层NN上进行前向和后向传播的内存代价，假设我们将k个中间结果存储在图中，并在子路径上进行前向和后向传播时递归地应用相同的策略，有以下递归公式：
g ( n ) = k + g ( n k + 1 ) ↓ g ( n ) = k l o g k + 1 ( n ) g(n) = k+g(\frac{n}{k+1}) \\ \downarrow \\ g(n) = k\ log_{k+1}(n) g(n)=k+g(k+1n)↓g(n)=k logk+1(n)
举一个特例，假设k = 1可以得到 g ( n ) = l o g 2 n g(n) = log_2n g(n)=log2n，这说明加入训练一个n层神经网络，所有实现都仅需 O ( n ) O(n) O(n)的内存用于特征映射，这将需要 O ( l o g 2 n ) O(log_2n) O(log2n)的正向成本，因此不太可能普遍使用。

3.3 Guideline for DL Frameworks

在Sec3证明了能与Sec2的系统优化结合起来后，这有助于其他DL 框架：

Enable options 去删除low cost ops产生的结果
提供planning alg.来提供有效的内存方案
使用户能够在computation graph中设置 mirror attribute 来进行内存优化

Search over Budget B

Alg.3使我们能在单个参数B的情况下生成优化的内存方案，这个算法依靠 approximate memory estimation 来获得更快的速度，在方案完成后，我们可以使用static allocation Alg.来计算准确需要的内存成本，之后我们能够在B上进行grid search以找到一个好的方案。

为了获得grid的设置，我们首先以B=0进行分配算法，然后改用 B = x y B = \sqrt{xy} B=xy 运行一次，这里的x和y是Alg.3在第一次运行中的输出。

x为存储inter-stage feature maps的近似成本
y为运行each stage产生的近似成本

B = x y B = \sqrt{xy} B=xy 为每一阶段内存成本的估计，然后围绕B= x y \sqrt{xy} xy 设置网格，进一步优化解决方案。

作者在实践中发现在 [ B 2 , 2 B ] [\frac{B}{\sqrt2},\sqrt{2}B] [2 B,2 B]上使用大小为6的网格可以给出实验中比较好的内存方案，MXNet用实现了分配算法，但没有尝试优化速度。

Answers 补充

S 1.1 – forward pass computation

如果用for循环一个样本一个样本的计算，显然比较慢，但是如果使用Vectorization，把m个样本压缩成一个向量X来计算，同样的把z、a都进行Vectorization处理得到Z、A，这样就可以对m的样本同时进行表示和计算。

提出原因

for iteration = 1 to 2000:  # 2000次迭代for i = 1 to m:           # 遍历m个样本for j = 1 to n          # 遍历每一个特征对应的w的梯度

时间复杂度 O ( n 3 ) O(n^3) O(n3)，太高了。可以通过Vectorization来消除第二个和第三个for循环，因为一个样本的n个特征可以组成一个向量，m个样本也可以组成一个大矩阵.

设置 X = [ x ( 1 ) , x ( 2 ) , . . . , x ( m ) ] m × n X = [x^{(1)},x^{(2)},...,x^{(m)}]_{m\times n} X=[x(1),x(2),...,x(m)]m×n,其中 x ( i ) x^{(i)} x(i)是一个列向量 x ( i ) = [ x 1 ( i ) , x 2 ( i ) , . . . , x n ( i ) ] T x^{(i)} = [x_{1}^{(i)},x_{2}^{(i)},...,x_{n}^{(i)}]^T x(i)=[x1(i),x2(i),...,xn(i)]T，包含了n个特征。

对应的n个权重也可以定义成一个列向量： W = [ w 1 , . . . , w n ] T W = [w_1,...,w_n]^T W=[w1,...,wn]T

将上面的算法都转换成矩阵或者向量:
Z = [ z ( 1 ) , . . . , z ( m ) ] = W T X + [ b , . . . , b ] = n p . d o t ( W T , X ) + b A = [ a ( 1 ) , . . . , a ( 1 ) ] = σ ( Z ) , Y = [ y ( 1 ) , . . . , y ( m ) ] d Z = [ d z ( 1 ) , . . . , d z ( m ) ] = A − Y d W = 1 m X d Z T = 1 m n p . d o t ( X , d Z . T ) d b = 1 m ∑ i = 1 m d z ( i ) = 1 m n p . s u m ( d Z ) W : = W − α d W b : = b − α d b \begin{aligned} Z &= [z^{(1)},...,z^{(m)}] = W^TX+[b,...,b] = np.dot(W^T,X)+b \\ A &= [a^{(1)},...,a^{(1)}] = \sigma(Z),\ Y =[y^{(1)},...,y^{(m)}]\\ dZ &= [dz^{(1)},...,dz^{(m)}] = A - Y \\ dW &= \frac{1}{m}XdZ^T = \frac{1}{m}np.dot(X,dZ.T) \\ db &= \frac{1}{m}\sum^{m}_{i=1}dz^{(i)} = \frac{1}{m}np.sum(dZ) \\ W &:= W-\alpha dW \\ b &:= b-\alpha db \end{aligned} ZAdZdWdbWb=[z(1),...,z(m)]=WTX+[b,...,b]=np.dot(WT,X)+b=[a(1),...,a(1)]=σ(Z), Y=[y(1),...,y(m)]=[dz(1),...,dz(m)]=A−Y=m1XdZT=m1np.dot(X,dZ.T)=m1i=1∑mdz(i)=m1np.sum(dZ):=W−αdW:=b−αdb
所谓的Vectorization，将Logistic Regression算法向量化的过程总结一下就是：

把m个样本，同时计算，同时算出它们的 z ( i ) z^{(i)} z(i)，即直接计算Z这个m维行向量
同时把Z的m维都激活，得到m维行向量A
得到A和Z之后，可以直接计算J对Z的梯度dZ了，得到dZ之后，可以直接算出W和b的梯度
同时更新所有的 w ( i ) w^{(i)} w(i)和b

用公式表示一下两层NN的前向传播过程：

Layer 1： Z [ 1 ] = W [ 1 ] ⋅ X + b [ 1 ] A [ 1 ] = σ ( Z [ 1 ] ) Z[1] = W[1]·X + b[1]A[1] = \sigma(Z[1]) Z[1]=W[1]⋅X+b[1]A[1]=σ(Z[1])

Layer 2: Z [ 2 ] = W [ 2 ] ⋅ A [ 1 ] + b [ 2 ] A [ 2 ] = σ ( Z [ 2 ] ) Z[2] = W[2]·A[1] + b[2]A[2] = \sigma(Z[2]) Z[2]=W[2]⋅A[1]+b[2]A[2]=σ(Z[2])

从上面可以知道，X其实就是A[0],不难看出，每一层的计算都是一样的，即

Layer i: Z [ i ] = W [ i ] ⋅ A [ i − 1 ] + b [ i ] A [ i ] = σ ( Z [ i ] ) Z[i] = W[i]·A[i-1]+b[i]A[i] = \sigma(Z[i]) Z[i]=W[i]⋅A[i−1]+b[i]A[i]=σ(Z[i])

对于损失函数，就跟Logistic regression中的一样，使用**“交叉熵（cross-entropy）”**，多分类问题： L = − ∑ y ( j ) ⋅ y j L = -\sum y(j)·y^{j} L=−∑y(j)⋅yj

以上是每个样本的loss，一般还需要计算整个样本集的loss，即cost，用J表示，J为L的平均：
J ( W , b ) = 1 m ⋅ ∑ ( y i , y ( i ) ) J(W,b) = \frac{1}{m·\sum (y^i,y(i))} J(W,b)=m⋅∑(yi,y(i))1
上面求Z、A、L、J的过程就是正向传播

S 2.1 Coarse-grained v.s Fine-grained

从概念理解

“粒度是根据项目模块划分的细致程度区分的，一个项目模块（或子模块）分得越多，每个模块（或子模块）越小，负责的工作越细，说明粒度越细，否则为粗粒度”

粗粒度和细粒度的区别主要出于重用的目的，像类的设计，为尽可能重用，采用细粒度的设计模式，将一个复杂的类（粗粒度）拆分为高度重用的职责清晰的类（细粒度）。

关于数据库的设计，原则上需要尽量减少表的数量及表与表之间的连接，能够设计成一个表的情况就不需要细分，所以可以考虑使用粗粒度的设计方式。

由上：

关于重用，粒度越粗越复杂，重用性应该是越差的；
关于数据库设计，粗细粒度的取舍也是一个关键

从接口角度理解

细粒度查询任务接口

interface TaskService{public List getTaskById(int id);public List getTaskByName(String name);public List getTaskByAge(int age);
}

粗粒度查询任务接口

//Person类有name，id及age属性
interface TaskService{public List getTask(Person person);
}

总结一下

coarse-grained 和 fine-grained是一个相对的概念，区别主要出于重用的目的。

类设计： 将复杂类（粗粒度）拆分为高度重用的职责清晰的类（细粒度）；
数据库设计： 尽量减少表的数量 & 表与表之间的连接，能够设置成一个表的情况就不需要细分；

数据库访问控制的粗细粒度问题：

根据控制对象的粗细程度，访问控制可以分为粗粒度和细粒度两种

粗粒度：通常为把规定访问整个数据库表或由基本表导出视图的某个层称为粗粒度的访问控制
细粒度：把安全控制细化到数据库的行级或列级