经典CNN模型-DenseNet的内存高效实现

  |  

摘要: 2018 年左右 DenseNet 内存高效实现的笔记

【对数据分析、人工智能、金融科技、风控服务感兴趣的同学,欢迎关注我哈,阅读更多原创文章】
我的网站:潮汐朝夕的生活实验室
我的公众号:潮汐朝夕
我的知乎:潮汐朝夕
我的github:FennelDumplings
我的leetcode:FennelDumplings


技术报告: DenseNet 的内存高效实现

在文章经典CNN模型-DenseNet手写笔记中,我们详细学习了 DenseNet 的论文,本文我们来看一下在深度学习框架中如何对 DenseNet 进行内存高效的实现。

原始的 DenseNet 很耗存储,本文给出了 Memory-Efficient 实现,论文如下:

空间爆炸的问题

DenseNet 训练时很耗内存。深度学习框架对稠密链接支持不好,所以只能借助反复 concatenate。对于大多数框架,每次拼接操作都会开辟新内存,来保存拼接后的特征(PyTorch、TensorFlow),这就导致 L 层的网络要小号相当于 L(L+1)/2 层的内存,第 I 层的输出存了 L - I + 1 份。

每层只产生 K 个 Featrue Map (K 很小,例如 12 ~ 48),但是用所有浅层得 feature 作为输入,使得参数量为 $O(网络深度^{2})$,这并不是问题,但是 DenseNet 的原始实现中存储使用量为 $O(FeatureMap数量^{2})$,因为中间 feature 来自各层 BN 和 Concatenate 后的输出,前向(计算下一个feature)和后向(求梯度)都要用,这个问题就很大了。

在实现中,可以在各层覆盖前一层的结果,被覆盖的值在反向时可以重新计算。这样 Feature Map 的空间消耗可以从平方降到线性,增加 15 ~ 20% 的训练时间。

分析一个 Dense Block (m 层)

最后一层的输入为 $[x_{1}, …, x_{m-1}]$。卷积特征的数量 ~ $O(网络深度)$ ~ $O(m)$。

但若网络也存储中间特征图,例如 BN 的输出,则存储空间可能扛不住,因为每层都计算了中间特征图,所以它们被存储产生了 $O(m^{2})$ 的内存使用。

而前面说的每层计算的中间特征图,很多框架将这些特征图保存在显存中,以便 BP 的过程使用。

卷积特征和参数的梯度通常是中间输出的函数,因此 BP 过程中间输出必须保持可访问。在 DenseNet 中,有两个操作导致了这种 $O(m^{2})$ 的内存使用:

(1) 预激活 BN

预激活对输入特征应用缩放和偏置。可产生高达每层 m 个归一化副本,因为每个副本有不同的缩放或偏置,因此框架上原始实现通常为这 $(m-1)(m-2)/2$ 个特征图分别分配内存。

(2) Concatenate

当所有输入在内存中连续的时候,conv 很快(最快),有的框架明确有这个要求,例如 PyTorch。CuDNN 没有要求,但不用连续内存会增加 30 ~ 50% 时间开销。

如果为了连续,每层必须将先前特征复制进连续内存。

GPU Conv 运算(如 CuDNN)假设小批量是最低维度,即:连续的内存中存的是小批量内不同张量中相同坐标的值,而非我们预期的 Feature Map 排列顺序(GPU 计算单元中一组 PE 的访存优化,参考 GPU 架构文档)。

原始实现

在框架中,计算操作是计算图的边,存储在内存的中间特征图是计算图的点。

step1: Input,前层特征因为不连续,所以复制并存入连续块,l 个前层,每个生成 k 个特征,这样 连续块必须能容纳 $l \times k$ 个特征图。
step2: 进入 BN,类似地分配 $l\times k$,ReLU 不分配,而是在现有内存上计算
step3: 卷积,以 BN 输出生成 k 个新特征。

两个中间操作需要 $O(lk)$,输出特征只需要每层一个恒定的 $O(k)$。

一些框架中,BP 会分配更多的内存:

(1) 用输出特征的梯度和从正向的 BN 特征图计算 BN 的梯度,需要额外 $O(lk)$ 存这些梯度。
(2) 类似地,需要 $O(lk)$ 哥而外内存,保存 Concatenate 的特征梯度。

节约内存的实现

由于 Concatenate 和 Normalization 这连个操作在计算上是比较便宜的。因此用两个预分配的共享内存池

  • FP: 保存所有中间输出
  • BP: 根据需要重算 Concatenate 和归一化的特征

这种重新计算的策略,已经被其他网络结构使用,例如《Training Deep Nets with Sublinear memory Cost》。

但注意,重计算中间输出,需要权衡时间换空间是否值得,对于 DenseNet 是有效的。

具体计算过程如下:

FP:concatenate 中间结果保存到共享内存1,需要 $O(lk)$;BN 中间结果保存到共享内存2,需要 $O(m)$。

BP:Concatenate、BN、Conv 在反向传播中都产生梯度张量。这个张量数据存在单个共享内存分配中,避免梯度的存储平方增长,PyTorch 中是这样实现的。

我们用类似于 Facebook 的 ResNet 中的共享内存方案。

参数量 ~ $O(深度^{2})$,由于存参数所需内存远小于存特征图所需内存,所以本文的方法在 DenseNet 中是有效的。

DenseNet 拾遗

CNN 提高效果的方向

  • 深:ResNet
  • 宽:Inception
  • 对 Feature 的极致利用:DenseNet

全文仅两个公式

第一个公式是 ResNet 中,值相加,通道数不变的公式。

第二个公式是 DenseNet 中,类似于 Inception 的通道合并。

这两个公式正是 ResNet 和 DenseNet 的本质区别。

DenseNet 和 Stochastic Depth 的关系

在 Stochastic Depth 中,Residual 中的 layers 在训练过程中会被随机 drop 掉,这就会使相邻层之间直接连接,这和 DenseNet 很像。

Bense Block 内部用 Bottleneck 层来减少计算量

因为深层后的输入会非常大,因此 Dense Block 内部可以用 Bottleneck 层来减少计算量,也及时在原有结构中增加 Conv1x1,即 ResNet-B 的【BN + ReLU + Conv1x1 + BN + ReLU + Conv3x3】。

Transition 层

连接两个相邻 Dense Block,降低特征图大小:【BN + ReLU + Conv1x1 + 2x2 avgpool】。

其中【BN + ReLU + Conv1x1】可以压缩模型,m 通道进 Transition,可以压缩成 $\theta m$。

Deep Supervision

输入层可以直达最后的误差信号,相当于隐式 Deep Supervision。

从哪些启发提出的 DenseNet

《Deep Networks with stachastic depth》中用了一种类似 Dropout 的方法改进 ResNet,训练中每一步都随机扔一些层,可以显著提升泛化性能。

这就带来两点启发:

  1. NN 不一定要是递进层级结构,即某一层不仅依赖上一层的特征,还可以依赖更前层的特征。
  2. 训练中随机扔很多层也不破坏收敛,说明 ResNet 有明显冗余



Share