DS-Net:可落地的动态网络,实际加速1.62倍,快改造起来 | CVPR 2021 Oral

论文提出能够适配硬件加速的动态网络DS-Net,通过提出的double-headed动态门控来实现动态路由。基于论文提出的高性能网络设计和IEB、SGS训练策略,仅用1/2-1/4的计算量就能达到静态SOTA网络性能,实际加速也有1.62倍

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

论文: Dynamic Slimmable Network

[图片上传失败...(image-262f83-1657684119973)]

Introduction


[图片上传失败...(image-37dd84-1657684119973)]

 模型速度在模型的移动端应用中十分重要,提高模型推理速度的方法有模型剪枝、权值量化、知识蒸馏、模型设计以及动态推理等。其中,动态推理根据输入调整其结构,降低整体计算耗时,包含动态深度和动态维度两个方向。如图2所示,动态网络自动在准确率和计算量之间trade-off,比静态的模型设计和剪枝方法要灵活。

[图片上传失败...(image-4ca994-1657684119973)]

 然而,论文发现包含动态维度的网络的实际运行速度大都不符合预期,主要原因在于动态剪枝后的稀疏卷积与当前硬件的计算加速不匹配。大多数卷积核的动态剪枝通过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


[图片上传失败...(image-986d99-1657684119973)]

 论文提出的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}),这样的卷积可表示为:

[图片上传失败...(image-328335-1657684119973)]

 动态卷积根据输入去掉不重要的特征通道,降低理论计算量,但其实际加速大都不符合预期。由于通道的稀疏性与硬件加速技术不匹配,在计算时不得不多次索引和拷贝需要的权值到新的连续内存空间再进行矩阵相乘。为了更好地加速,卷积核在动态权值选择时必须保持连续且相对静态。
 基于上面的分析,论文设计了结构路由器\mathcal{A}(\theta),能够偏向于输出稠密的选择结果。对于N输出、M输入的卷积核W\in\mathbb{R}^{N\times M},结构路由器输出精简比例\rho\in(0,1],通过切片操作[:]选择卷积核的前\rho\times N部分构成切片动态卷积:

[图片上传失败...(image-290425-1657684119973)]

[:]切片操作加*稠密矩阵乘法要比索引操作或稀疏矩阵相乘要高效得多,保证了实际运行时的速度。

  • 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^{'}为在线网络和目标网络:

[图片上传失败...(image-d96997-1657684119973)]

\alpha为动量因子,控制历史参数的比例,t为训练轮次。在训练时,模型的EMA会比在线网络更加稳定和准确,为精简子网提供高质量的训练目标。
 接着,参考MealV2使用一组teacher网络来生成更多样的输出向量供student网络学习的做法,在进行in-place distillation时使用不同的子网构成一组teacher网络,主要提供目标向量给最小子网学习。

[图片上传失败...(image-e7640e-1657684119973)]

 整体训练过程如图4所示。结合sandwich rule和上述优化的in-place distillation,每论训练有以下3种网络:

  • 最大的子网L使用数据集标签作为训练目标。
  • n个随机维度的子网使用目标网络的最大子网的向量输出作为训练目标。
  • 最小的子网使用上述子网在目标网络中对应的子网的向量输出的组合作为训练目标,即训练目标为:

[图片上传失败...(image-19eaac-1657684119973)]

 总结起来,超网训练的IEB损失为:

[图片上传失败...(image-88f56e-1657684119973)]

Dynamic Slimming Gate

[图片上传失败...(image-2b0fbd-1657684119973)]

 这里先介绍公式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})}转换为两个函数的组合:

[图片上传失败...(image-f2c7aa-1657684119973)]

\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为中间特征维度,gL_p的长度):

[图片上传失败...(image-f5a9d2-1657684119973)]

 以图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的点积得到预测的精简比例:

[图片上传失败...(image-4ad19c-1657684119973)]

 论文采用的精简比例生成方法跟通道注意力方法十分类似,通过添加第三个全连接层W_3^{\rho_{n-1}\times d},可直接为网络引入注意力机制。基于上面的结构,论文提出double-headed dynamic gate,包含用于通道路由的hard channel slimming head以及用于通道注意力的soft channel attention head,其中soft channel attention head定义为:

[图片上传失败...(image-3e8477-1657684119973)]

\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损失定义为:

[图片上传失败...(image-a1e1f1-1657684119973)]

\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


[图片上传失败...(image-61777d-1657684119973)]

[图片上传失败...(image-fdfe96-1657684119973)]

 与不同类型的网络对比ImageNet性能。

[图片上传失败...(image-f6df86-1657684119973)]

 CIFAR-10性能对比。

[图片上传失败...(image-24503c-1657684119973)]

 VOC检测性能对比。

[图片上传失败...(image-596caf-1657684119973)]

 对IEB训练方法各模块进行对比实验。

[图片上传失败...(image-8e931d-1657684119973)]

 对比SGS损失与精简比例分布的可视化。

[图片上传失败...(image-7a92da-1657684119973)]

 对比不同的SGS训练策略,Try Best为本文的策略,Give up为放弃hard samples,将其归类为最小精简网络的目标。

[图片上传失败...(image-3b9336-1657684119973)]

 对比不同门控设计细节。

Conclusion


 论文提出能够适配硬件加速的动态网络DS-Net,通过提出的double-headed动态门控来实现动态路由。基于论文提出的高性能网络设计和IEB、SGS训练策略,仅用1/2-1/4的计算量就能达到静态SOTA网络性能,实际加速也有1.62倍。



如果本文对你有帮助,麻烦点个赞或在看呗~
更多内容请关注 微信公众号【晓飞的算法工程笔记】

work-life balance.
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 194,670评论 5 460
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 81,928评论 2 371
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 141,926评论 0 320
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 52,238评论 1 263
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 61,112评论 4 356
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 46,138评论 1 272
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 36,545评论 3 381
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 35,232评论 0 253
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 39,496评论 1 290
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 34,596评论 2 310
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 36,369评论 1 326
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 32,226评论 3 313
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 37,600评论 3 299
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 28,906评论 0 17
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 30,185评论 1 250
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 41,516评论 2 341
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 40,721评论 2 335

推荐阅读更多精彩内容