这篇文章是PPOPP'18的文章，感觉在该方向的共享度比较大，从16年的vDNN开始，到18年该论文再到20年ASPLOS的Capuchin，方案每两年前进一步。这里将这个文章详细总结一下，方便日后阅读。
一、前言
背景问题还是老生常谈，即DNN在训练的过程中GPU不够用。
文章提出现在的框架中也存在一定的内存优化方案，但是这些方案都有一些缺陷，例如：
caffe与Torch对非线性的网络支持效果不好
MXNet忽略了网络层之间的内存差异
该文章基于tensor为粒度，设计了三个方案，从而一步一步降低内存的使用率，具体的方案后续逐步讲解。
本文有如下贡献：
1 阐述了非线性网络的挑战，并讨论了现在方案的不足；
2 为卷积空间动态分配内存，从而提升系统性能；
3 设计并实现了SuperNeurons，使得网络能够超越gpu内存限制，从而更高效；
二、背景
这里简单的交代一些文章中总结的比较好的背景知识。
传统的卷积神经网络中包括如下类型的层：
Convolution (CONV), Pooling (POOL), Activation (ACT), Softmax, Fully Connected (FC), Local Response Normalization (LRN), Batch Normalization (BN), and Dropout
然而为了解决传统线性神经网络结构造成的梯度消失问题，非线性结构出现，其中包括fan类型与join类型：
这些非线性结构网络通常网络非常深，计算非常密集从而训练过程比较占内存，且时间长。
第一个挑战是GPU内存有限。
神经网络中有快速与慢速的卷积函数，而快速的卷积函数需要更多的工作空间。
上图中显示了快速与慢速卷积函数对应的空间使用以及效率提升情况。
从图中也能够看出非线性的网络内存占用量是非常巨大的，传统的gpu内存装不下这么大的网络。
第二个挑战是非线性网络层与层之间的依赖是变化的。
这种结构使得网络优化方案的设计要更具有动态性。
许多框架使用静态方案来优化内存：
Caffe与Torch选择在反向传播的时候直接使用前向传播的数据，从而避免重新生成的过程。不过这种方案需要提前知道未来的训练依赖；
MXNet与TF使用DAG引擎，DAG的好处是用户能够直接定义出层与层之间的依赖，从而使得系统有一个预先的知识。从而能够将不需要的tensor释放掉。
MXNet能够在前传的时候将部分tensor释放掉，但是这种方法忽略了网络层的不均匀内存分配，因此需要大量不必要的内存使用。
TF将寿命长的数据交换出去，但是并没有使用一些优化方案，比如pinned data transfer。（先传到cache，再传到gpu。这样太慢，pinned可以直接传到gpu）
三、设计思路
这里定义Baseline的最大内存使用量为：
1 ~ N的前传内存 + 1 ~ N的反传内存
此外定义网络层中的layer最大内存占用量为：
其中i属于1~N。
- 1 Liveness Analysis
该方案将Baseline的最大内存使用量减少为：
- 2 Unified Tensor Pool(UTP)
第二种方案将内存使用量进一步降低到：
其中Checkpoint代表网络中的全连接层与卷积层。
- 3 Cost-Aware Recomputation
第三种方案将内存使用量进一步降低到：
下面来看具体的方案：
为了对非线性结构的网络有一个宏观的概念，论文首先对该网络进行了一个探测。
template <class value_type>
void network_t<value_type>::fsetup_kernel(base_layer_t<value_type>* b) {
if(b == NULL) return;
//conduct forward computation this->fcounter = this->fcounter + 1;
b->fcounter_inc();
// b represents i in paper, the size of i is three(cause there are three forks join into i),
// but until fcounter equals 3 that 'i' can go on executing
// that is why >= is ok
if(b->get_fcounter() < b->get_prev_size() ) {
return;
}
b->forward_setup(reg, &cudnn_handle);
this->reg->register_net_layers(b->get_base_id(), (void*) b);
this->reg->register_net_comp_route(b->get_base_id(), FORWARD);
std::vector<base_layer_t<value_type>*> next = b->get_next();
// next.size is the fork number. b's next.size() is 2(one for 'c' and one for 'd')
if(next.size() == 1) {
//regular network layer
fsetup_kernel(next[0]);
} else if(next.size() > 1) {
//fork layer
for(size_t i = 1; i < next.size(); i++) {
// execute one by one
fsetup_kernel(next[i]);
}
fsetup_kernel(next[0]);
}
// this->fcounter = 0;
b->reset_fc_counter();
}
具体代码如上所述，流程大约是1~4 。通过深度优先的方案获取到非线性网络的执行顺序。
1 Liveness Analysis
思想是维护一个in与out表，其中in代表当前层计算所需tensor，而out是输出的live的tensor，其中的in-out差值就是释放掉的无用tensor。
所以这个表的简历需要逐层遍历网络，从第一个开始N−1, 2 N − 2, ..., 2, 1。
简单来说就是，后面的层不需要的tensor就释放掉。
由于前向传播的数据会被后项传播使用，所以该方案的最大内存使用量为：前传所有tensor量+第N层的后项传播使用量（之后反向传播操作就可以利用前面的无用空间，从而节约了内存）
第k步反向传播内存使用量为：
所以该步骤之后，最大的内存占用量为：
即在最后一步N，内存达到最大。之后便进入反向传播。
然而上述方法由于释放申请内存操作过于频繁从而浪费了大量的时间，于是系统预分配了一个大的gpu内存池，从池子里找空闲位置更快。
2 Unified Tensor Pool(UTP) and
该方案就是常见的GPU与DRAM的交换方案。
不是所有的层都适合交换，所以根据下图，我们仅选择卷积层进行交换。
原因是：
POOL, ACT, BN and LRN all together occupy over 50% of the total memory, while their computations only account for an average of 20% of the entire workload.
而It is also not fruitful to offload on Dropout, Softmax and FC layers since they only use less than 1% of the total memory内存占用太小，不值当。
Offloading转移出去：运行时异步将卷积层的数据转移到DRAM中，并使用pinned方案。注意这里的中间过程时间是当前
卷积层——下一个卷积层。Prefetching预取：反向传播时在上一个卷积操作时预取下一个卷积所需数据。
该方案进一步降低内存使用率到（卷积层的数据释放掉了）：
如果简单的将 所有计算过并且对当前层计算无用的卷积层输出 暂存到CPU中，会涉及许多设备间通信，极有可能拖慢训练速度。论文中采用了类似缓存的思想，定义了Caching tensors on GPU DRAM。即只有在空间不足的情况下才将tensor暂存，选取暂存tensor的策略为Least Recent Used (LRU)，即最长时间未被使用的tensor将被暂存到CPU内存中。论文中使用LRU的motivation是BP算法中最新使用的tensor最早用于梯度值计算。
这里思想确实有点粗粒度，毕竟如果当前内存量足够，但是后面不够了同样会使得程序调用失败
3 Cost-Aware Recomputation
对于POOL, ACT, LRN and BN等层来说，他们使用了50%的内存但是计算却只用了10%。
所以这些层可以释放掉并重计算回来。
在每个区块内，存在两种重计算的策略：speed-centric和memory -centric，即速度优先和显存优先。速度优先存储每一层重计算的结果，直至不再需要或者区块结束，从而仅需要对每一层进行一次重计算，计算复杂度为O(N)，多存储的tensor数为NN（其中NN为该recomputation segment内的神经层数）。memory优先则是每次仅存储重计算的输出，而不存储中间结果；若中间结果在后续计算中被依赖，则在需求时重新计算；其计算复杂度为O(N2)。
简单来说，A-B-C-D如果从A计算D可以同时把BC计算出来，如果内存够就把他们留着，不够就不留。比较动态。
所以理论最优内存量是当前层的所需的内存。
四、实验
实验的话就不进行详细总结了，放一个总图：
