论文提出可能适配硬件加速的动静网络 DS-Net,通过提出的 double-headed 动静门控来实现动静路由。基于论文提出的高性能网络设计和 IEB、SGS 训练策略,仅用 1 /2-1/ 4 的计算量就能达到动态 SOTA 网络性能,理论减速也有 1.62 倍
起源:晓飞的算法工程笔记 公众号
论文: Dynamic Slimmable Network
- 论文地址:https://arxiv.org/abs/2103.13258
- 论文代码:https://github.com/changlin31/DS-Net
Introduction
模型速度在模型的挪动端利用中非常重要,进步模型推理速度的办法有模型剪枝、权值量化、常识蒸馏、模型设计以及动静推理等。其中,动静推理依据输出调整其构造,升高整体计算耗时,蕴含动静深度和动静维度两个方向。如图 2 所示,动静网络主动在准确率和计算量之间 trade-off,比动态的模型设计和剪枝办法要灵便。
然而,论文发现蕴含动静维度的网络的理论运行速度大都不合乎预期,次要起因在于动静剪枝后的稠密卷积与以后硬件的计算减速不匹配。大多数卷积核的动静剪枝通过 zero masking(惯例卷积后再通过 mask 取对应的输入)或 path indexing(间接通过 $[:,:]$ 获取对应的新卷积再计算)来实现,如表 1 所示,这些办法的计算效率都不高,导致整体推理速度没有放慢。
为了解决这一问题,论文提出了动静可精简网络 DS-Net,在实现动静网络的同时也有很好的硬件匹配性。
论文的次要奉献如下:
- 提出新的动静网络路由机制,通过提出的 double-headed 动静门控来实现网络结构的动静路由。另外,卷积的动静剪枝通过切片的形式放弃权值的内存连续性,能够很好地适配硬件加速。
- 提出用于 DS-Net 的两阶段训练形式,蕴含 IEB 和 SGS 办法。IEB 用于稳固可精简网络的训练,SGS 用于进步门控输入的多样性,两者都能帮忙进步 DS-Net 的性能。
- 通过 ImageNet 试验比照,DS-Net 的整体性能比 SOTA 动静网络高约 5.9%,比 ResNet 和 MobileNet 等动态网络性能略微降落,然而有 2 - 4 倍计算量节俭以及 1.62 倍理论推理减速。
Dynamic Slimmable Network
论文提出的 dynamic slimmable network 通过学习可精简的超网 (supernet) 以及动静门控 (gating) 机制,达到依据不同输出样本动静生成网络的目标。如图 3 所示,DS-Net 的超网为蕴含全副残缺卷积的残缺网络。动静门控则是一系列预测模块,依据输出动静设定每个阶段的卷积维度,进而生成子网,这一过程也称为动静路由(dynamic routing)。
目前的动静网络钻研中,主网络和动静路由通常是联结训练的,相似于联结优化的网络搜寻办法。参考 one-shot NAS 办法,论文提出解耦的两阶段训练方法来保障 DS-Net 中每个门路的泛化性。在 stage I 中,禁用门控的性能并用 IEB 办法训练超网,在 stage II 中,固定超网的权值独自用 SGS 办法训练门控。
Dynamic Supernet
这里先介绍可在硬件高效运行的通道切片办法以及论文设计超网,而后再介绍 Stage I 中用到的 IEB 办法。
-
Supernet and Dynamic Channel Slicing
在如动静裁剪、动静卷积等动静网络中,卷积核 $\mathcal{W}$ 依据输出 $\mathcal{X}$ 进行动静参数化 $\mathcal{A}(\theta, \mathcal{X})$,这样的卷积可示意为:
动静卷积依据输出去掉不重要的特色通道,升高实践计算量,但其理论减速大都不合乎预期。因为通道的稠密性与硬件加速技术不匹配,在计算时不得不屡次索引和拷贝须要的权值到新的间断内存空间再进行矩阵相乘。为了更好地减速,卷积核在动静权值抉择时必须放弃间断且绝对动态。
基于下面的剖析,论文设计了构造路由器 $\mathcal{A}(\theta)$,可能偏差于输入浓密的抉择后果。对于 $N$ 输入、$M$ 输出的卷积核 $W\in\mathbb{R}^{N\times M}$,构造路由器输入精简比例 $\rho\in(0,1]$,通过切片操作 $[:]$ 抉择卷积核的前 $\rho\times N$ 局部形成切片动静卷积:
$[:]$ 切片操作加 $*$ 浓密矩阵乘法要比索引操作或稠密矩阵相乘要高效得多,保障了理论运行时的速度。
-
SuperNet
将多个动静卷积组合起来即可搭建超网,超网通过设置不同的特色维度组合创立多个子网。将构造路由器禁用时,超网等同于常见可精简网络,可用相似的办法进行预训练。
-
In-place Ensemble Bootstrapping
经典的 Universally Slimmable Networks 通过两个办法来无效地晋升整体的性能:
- sandwich rule:每次训练的网络组合蕴含最大的子网、最小的子网以及其它子网,其中最大的子网和最小的子网别离决定了可精简网络性能的上界和下界。
- in-plcae distillation:将最大子网的向量输入作为其它子网的训练指标,而最大子网的训练指标则是数据集标签,这样对可精简网络更好地收敛有很好的帮忙。
尽管 in-place distillation 很无效,但最大子网权值的激烈抖动会导致训练难以收敛。依据 BigNas 的试验,应用 in-place distillation 训练较为简单的网络会极其不稳固。如果没有残差连贯或非凡的权值初始化,在训练初期甚至会呈现梯度爆炸的状况。为了解决可精简网络收敛难的问题并且晋升整体性能,论文提出了 In-plcae Ensemble Boostrapping(IEB)办法。
首先,参考 BYOL 等自监督和半监督办法,应用过往的表达能力进行自监督的 in-plcae distillation 训练的做法,将模型的指数滑动均匀 (EMA, exponential moving average) 作为指标网络生成指标向量。定义 $\theta$ 和 $\theta^{‘}$ 为在线网络和指标网络:
$\alpha$ 为动量因子,管制历史参数的比例,$t$ 为训练轮次。在训练时,模型的 EMA 会比在线网络更加稳固和精确,为精简子网提供高质量的训练指标。
接着,参考 MealV2 应用一组 teacher 网络来生成更多样的输入向量供 student 网络学习的做法,在进行 in-place distillation 时应用不同的子网形成一组 teacher 网络,次要提供指标向量给最小子网学习。
整体训练过程如图 4 所示。联合 sandwich rule 和上述优化的 in-place distillation,每论训练有以下 3 种网络:
- 最大的子网 $L$ 应用数据集标签作为训练指标。
- $n$ 个随机维度的子网应用指标网络的最大子网的向量输入作为训练指标。
- 最小的子网应用上述子网在指标网络中对应的子网的向量输入的组合作为训练指标,即训练指标为:
总结起来,超网训练的 IEB 损失为:
Dynamic Slimming Gate
这里先介绍公式 2 中输入 $\rho$ 因子的构造路由器 $\mathcal{A}(\theta, \mathcal{X})$ 以及动静门控的 double-headed 设计,最初再介绍 Stage II 训练应用的 sandwich gate sparsification(SGS)办法。
-
Double-headed Design
将特色图转换为精简比例 $\rho$ 有两种办法:1)标量模式:间接通过 sigmoid 输入 0 到 1 的标量作为精简比例。2)one-hot 模式:通过 argmax/softmax 失去 one-hot 向量,抉择离散的候选向量 $L_p$ 中对应的精简比例。
论文对这两种办法进行比照后,抉择了性能更好的 one-hot 模式。为了将特色图 $\mathcal{X}$ 转换为 one-hot 向量,将 $\mathcal{A(\theta, \mathcal{X})}$ 转换为两个函数的组合:
$\mathcal{E}$ 将特色图下采样为向量,$\mathcal{F}$ 将向量转化为 one-hot 向量用于后续的维度切片。参考 DenseNet 等网络,$\mathcal{E}$ 为全局池化层,$\mathcal{F}$ 为全连贯层 $W_1\in\mathbb{R}^{d\times C_n}$+ReLU+$W_2\in\mathbb{R}^{g\times d}$+argmax 函数($d$ 为两头特色维度,$g$ 为 $L_p$ 的长度):
以图 3 的第 $n$ 个门控为例,将大小为 $\rho_{n-1}C_n\times H_n\times W_n$ 的特色图 $\mathcal{X}$ 转换成向量 $\mathcal{X}_{\mathcal{E}}\in \mathbb{R}^{\rho_{n-1}C_n}$,随后用 argmax 将向量进一步转换成 one-hot 向量,最初通过计算 one-hot 向量与 $L_p$ 的点积失去预测的精简比例:
论文采纳的精简比例生成办法跟通道注意力办法非常相似,通过增加第三个全连贯层 $W_3^{\rho_{n-1}\times d}$,可间接为网络引入注意力机制。基于下面的构造,论文提出 double-headed dynamic gate,蕴含用于通道路由的 hard channel slimming head 以及用于通道注意力的 soft channel attention head,其中 soft channel attention head 定义为:
$\delta(x)=1+tanh(x)$,channel attention head 参加 stage I 的训练。
-
Sandwich Gate Sparsification
在 stage II 训练中,论文应用分类穿插熵损失 $L_{cls}$ 和复杂度惩办函数 $L_{cplx}$ 来端到端地训练门控,疏导门控为每个输出图片抉择最高效的子网。为了可能用 $L_{cls}$ 来训练不可微的 slimming head,论文尝试了经典的 gumbel-softmax 办法,但在试验中发现门控很容易收敛到动态的选项,即便加了 Gumbel 噪声也优化不了。
为了解决收敛问题并且减少门控的多样性,论文提出 Sandwich Gate Sparsification(SGS)训练方法,应用最大子网和最小子网辨认输出图片中的 hard 和 easy,为其生成 slimming head 输入精简因子的 GT。基于训练好的超网,将输出大抵地分为三个级别: - Easy samples $\mathcal{X}_{easy}$:可能被最小子网辨认的输出。
- Hard samples $\mathcal{X}_{hard}$:不能被最大子网辨认的输出。
- Dependent samples $\mathcal{X}_{dep}$:不属于上述两种的输出。
为了最小化计算耗费,easy samples 应该都应用最小子网进行辨认,即门控的 GT 为 $\mathcal{T}(\mathcal{X}_{easy})=[1,0,\cdots,0]$。而对于 dependent samples 和 hard samples 则应该激励其尽量应用最大的子网进行辨认,即门控的 GT 为 $\mathcal{T}(\mathcal{X}_{hard})=\mathcal{T}(\mathcal{X}_{dep})=[0,0,\cdots,1]$。基于这些生成的门控 GT,SGS 损失定义为:
$\mathbb{T}_{sim}(\mathcal{X})\in{0,1}$ 代表 $\mathcal{X}$ 是否应该被最小子网预测,$\mathcal{L}_{CE}(\mathcal{X},\mathcal{T})=-\sum\mathcal{T}*log(\mathcal{X})$ 为门控输入与生成 GT 之间穿插熵损失。
Experiment
与不同类型的网络比照 ImageNet 性能。
CIFAR-10 性能比照。
VOC 检测性能比照。
对 IEB 训练方法各模块进行比照试验。
比照 SGS 损失与精简比例散布的可视化。
比照不同的 SGS 训练策略,Try Best 为本文的策略,Give up 为放弃 hard samples,将其归类为最小精简网络的指标。
比照不同门控设计细节。
Conclusion
论文提出可能适配硬件加速的动静网络 DS-Net,通过提出的 double-headed 动静门控来实现动静路由。基于论文提出的高性能网络设计和 IEB、SGS 训练策略,仅用 1 /2-1/ 4 的计算量就能达到动态 SOTA 网络性能,理论减速也有 1.62 倍。
如果本文对你有帮忙,麻烦点个赞或在看呗~
更多内容请关注 微信公众号【晓飞的算法工程笔记】