效率优先的ShuffleNet

这个shuffle并不是你想的那种shuffle。

效率优先的ShuffleNet

封面图片:Martin Adams on Unsplash


本文是对文章“ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices”的摘录。原文作者来自旷视科技。常规网络架构通过剪枝、压缩和量化等手段来加速网络。ShuffleNet则另辟蹊径、针对低运算能力设备设计基本网络结构。文中介绍了一种专为预算能力在10-150 MFLOPs移动设备设计的网络结构,在ARM架构下的运算速度相较于AlexNet提升了约13倍。你可以在这里找到原文:

ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices
We introduce an extremely computation-efficient CNN architecture namedShuffleNet, which is designed specially for mobile devices with very limitedcomputing power (e.g., 10-150 MFLOPs). The new architecture utilizes two newoperations, pointwise group convolution and channel shuffle, to greatly red…

一些先进网络架构在体量较小时,密集的1×1卷积拖累了网络的运算效率。使用pointwise分组卷积可以降低1×1卷积的运算复杂度。为了解决分组卷积的副作用,作者提出了通道重排来协助跨通道的信息流通,并在此基础上构建了ShuffleNet。在同等运算量的前提下,ShuffleNet支持更多的通道,这有助于小体量网络的表现。例如ImageNet上top-1错误率比MobileNet低7.8%;在ARM上速度比AlexNet快13倍。

过往研究

ShuffleNet所采用设计元素的相关研究。

分组卷积

分组卷积最早被应用在AlexNet中,目的是实现模型在两个GPU上的分布式计算(大显存的小伙伴这里可以偷笑了)。其有效性也在ResNeXt中得到验证。Xception提出的depthwise separable convolution则派生自Inception系列中的separable convolution,并在MobileNet中发扬光大。ShuffleNet以新的形式再利用了这两种操作。

通道重排

该操作并非首创,但是相关研究中它的目的与ShuffleNet存在差异没有被深入研究。

ShuffleNet的优化措施

可以概括为3点。

针对群组卷积的通道重排

Xception与ResNeXt并没有完全考虑1×1卷积,导致网络较小时通道数量受限。为解决这一问题,可以为1×1卷积添加分组卷积操作。然而该操作存在副作用:输出仅仅与输入通道的一小部分有关联。如下图(a)所示。

通道重排操作。作者:Xiangyu Zhang, et, al.

群组卷积可以尝试从不同的群组获得输入数据来解决这个问题,如上图(b)所示。进一步,可以通过通道重排来实现,如图(c)所示。假设卷积层分组数为$g$,输出通道数为$g\times n$;首先将输出通道的尺度reshape为$(g, n)$,之后转置、降维度并作为下一层的输入。该操作在两个卷积的组数不同时依旧有效并且可导。这个过程形似洗牌(shuffle),只不过是非随机的、有规律的shuffle。

ShuffleNet单元

在下图(a)的基础上,将1×1卷积替换为1×1分组卷积,且在第一个卷积后添加通道重排操作如(b)所示。第二个1×1分组卷积的目的在于恢复通道尺度以匹配捷径尺度,且该分组卷积后没有通道重排,因为作者发现影响不大。另外batch norm与ReLU的使用也有不同。

ShuffleNet单元。作者:Xiangyu Zhang, et, al.

针对降采样功能,首先在捷径路径添加stride为2的3×3 average pooling;然后将最后的求和操作改为拼接。

从FLOPs的角度考虑,假设输入尺度为 $c\times h \times w$,bottleneck的通道数为 $m$,ResNet需要$hw(2cm+9m^2)$ FLOPs,ResNeXt需要$hw(2cm+9m^2/g)$ FLOPs,而ShuffleNet需要 $hw(2cm/g+9m)$ FLOPs。其中$g$为分组数。换言之,ShuffleNet可以使用更多的通道数,这对小体量网络来说很重要。ShuffleNet的depthwise convolution仅仅在bottleneck特征图上执行,这是因为在低功耗设备上受内存访问成本的限制表现并不好。

网络架构

ShuffleNet的网络架构如下图所示,在3个stage内堆叠ShuffleNet单元。每个stage的第一个单元通过stride=2来降采样;stage最后输出的通道数翻倍。每个单元的bottleneck通道数为输出通道数的1/4。

ShuffleNet架构。作者:Xiangyu Zhang, et, al.

每个单元内,群组数$g$控制着pointwise卷积的稀疏程度。平衡$g$与输出通道数可以维持运算量在140 MFLOPs左右。

实验

作者在ImageNet 2012数据上开展了分类实验。Weight decay为4e-5,在数据增强时使用较小的缩放比例,因为小体量的网络更多的存在拟合不足而非过拟合。使用4 GPU训练3×105 步需要1到2天时间,batch size为1024。

实验发现:

  1. 分组卷积带来了更好的模型表现。
  2. 过多的分组会造成每个分组的输入变少,不利于模型表现。
  3. 通道重排能够显著提升模型表现。

其它实验内容涉及到与其它类别网络的比较,这里不再列出。


总结

ShuffleNet的核心思想在于通过通道重排改善分组卷积造成的跨组信息流通障碍。如果希望通过分组卷积来提速网络,该设计理念值得参考应用。