简单的特征值梯度剪枝，CPU和ARM上带来4-5倍的训练加速 | ECCV 2020

论文通过DBTD方法计算过滤阈值，再结合随机剪枝算法对特征值梯度进行裁剪，稀疏化特征值梯度，能够降低回传阶段的计算量，在CPU和ARM上的训练分别有3.99倍和5.92倍的加速效果

 

来源：晓飞的算法工程笔记 公众号

论文: Accelerating CNN Training by Pruning

Activation Gradients

• 论文地址：https://arxiv.org/abs/1908.00173

Introduction

---

  在训练过程中，特征值梯度的回传和权值梯度的计算占了大部分的计算消耗。由于这两个操作都是以特征值梯度作为输入，而且零梯度不会占用计算资源，所以稀疏化特征值梯度可以降低回传阶段的计算消耗以及内存消耗。论文的目标在于高效地降低训练负载，从而在资源有限的平台进行大规模数据集的训练。

  论文假设特征值梯度服从正态分布，基于此计算阈值$\tau$，随后使用随机剪枝算法(stochastic pruning)将小于阈值的特征值梯度随机置为零或$\pm \tau$。经理论推理和实验证明，这种方法不仅能够有效地稀疏化特征值梯度，还能在加速训练的同时，不影响训练的收敛性。

General Dataflow

---

  卷积层通常包含4个阶段：推理、特征值梯度回传、权值梯度计算和权值更新。为了表示这些阶段的计算，论文定义了一些符号：

  卷积层的四个训练阶段的总结为：

  论文通过可视化发现，回传阶段的特征值梯度几乎全是非常小的、接近于零的值，自然而然地想到将这些值去掉不会对权值更新阶段造成很大的影响，所以论文认为剪枝特征值梯度能够加速卷积层在训练时的计算。

Sparsification Algorithms

---

Distribution Based Threshold Determination (DBTD)

  剪枝操作最关键的步骤是决定选择哪些元素进行消除，先前有研究使用最小堆进行元素选择，但这会带来较大的额外计算开销。为此，论文采用简单的阈值过滤进行元素选择。

  论文首先分析了两种经典的卷积网络结构的特征值梯度分布：Conv-ReLU结构和Conv-BN-ReLU结构：

• 对于Conv-ReLU结构，输出的特征值梯度$dO$是稀疏的，但其分布是无规律的，而结构的输入特征值梯度$dI$几乎全是非零值。通过统计发现，$dI(\cdot)$的分布以零值对称分布，且密度随着梯度值的增加而下降。
• 对于Conv-BN-ReLU结构，BN层设置在卷积层与ReLU层中间，改变了梯度的分布，且$dO$的分布与$dI$类似，。

  所以，上述的两种结构的梯度都可认为服从零均值、方差为$\sigma^2$的正态分布。对于Conv-ReLu结构，由于ReLU不会降低稀疏性，$dO$能够继承$dI$的稀疏性，将$dI$是作为Conv-ReLU结构中的剪枝目标梯度$g$。而对于Conv-BN-ReLU结构，则将$dO$作为剪枝目标$g$。这样，两种结构的剪枝目标都可统一为正态分布。假设$g$的数量为$n$，可以计算梯度的绝对值的均值，并得到该均值的期望为：

  这里的期望为从分布中采样$n$个点的期望，而非分布的整体期望，再定义以下公式

  将公式2代入公式1中，可以得到：

  从公式3可以看出$\tilde{\sigma}$为参数$\sigma$的无偏估计，接近于真实的均值，且$\tilde{\sigma}$的整体计算消耗是可以接受的。基于上面的分析，论文结合正态分布的累积函数$\Phi$、剪枝率$p$和$\tilde{\sigma}$计算阈值$\tau$：

Stochastic Pruning

  剪枝少量值较小的梯度几乎对权值的更新没有影响，但如果将这些值较小的梯度全部设为零，则会对特征值梯度的分布影响很大，进而影响梯度更新，造成严重的精度损失。参考Stochastic Rounding算法，论文采用随机剪枝来解决这个问题。

  随机剪枝逻辑如算法1所示，对于小于阈值$\tau$的梯度值，随机采样一个缩放权重来计算新阈值，再根据新阈值将梯度值置为零或$\pm \tau$。

  随机剪枝的效果如图2所示，能够在保持梯度分布的数学期望的情况下进行剪枝，与当前的方法相比，论文提出的方法的优点如下：

• Lower runtime cost：DBTD的计算复杂度$O(n)$小于top-k算法$O(nlogk)$，且DBTD对硬件更友好，能够在异构平台实现。
• Lower memory footprint：随机裁剪能保持收敛性，且不需要存储而外的内存。

  至此，Sparsification Algorithms在梯度回传时的特征值梯度计算为：

Experimental Results

---

  在CIFAR-10、CIFAR-100以及ImageNet上进行准确率验证。

  在CIFAR-10和ImageNet上进行收敛性验证。

  在不同的设备上进行加速效果验证。

Conclustion

---

  论文通过DBTD方法计算过滤阈值，再结合随机剪枝算法对特征值梯度进行裁剪，稀疏化特征值梯度，能够降低回传阶段的计算量，在CPU和ARM上的训练分别有3.99倍和5.92倍的加速效果。论文提出的特征值稀疏化算法看似很简单，其实进行了充分的理论推导以及实验验证，才得到最终合理的过滤方法，唯一可惜的是没在GPU设备上进行实验验证。论文对算法的收敛性以及期望有详细的理论验证，不过这里没有列出来，有兴趣的可以去看看原文。

 

 

 

如果本文对你有帮助，麻烦点个赞或在看呗～

更多内容请关注 微信公众号【晓飞的算法工程笔记】