AI模型计算机视觉

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

字数统计: 1.8k字   |   阅读时长: 6分
, , ,

摘要: 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