ShuffleNet——面向移动设备的极为高效的卷积神经网络
1. 引言
上周我们介绍了 MobileNet,具有计算量小、参数量少的特点,适合在移动设备上部署。其网络结构特点是用可分离的深度方向卷积替代经典卷积计算。我们也注意到在 MobileNet 中计算量和参数数量大部分集中在 1x1 卷积部分(计算量 > 90%,参数量 > 70%),很自然的一个想法,能否进一步降低这部分计算量和参数量?
2. 模型简介
Face++ 最近提出的 ShuffleNet【1】给了一个很好的思路。该模型使用分组卷积(group convolution)、通道重排(channel shuffle)两项操作,在保持原有准确率的条件下大大降低计算量。
2.1 分组卷积
分组卷积概念最早在 AlexNet【2】中引入,用来将模型分到两块 GPU 上(当时由于单块 GPU GTX 580 显存只有 3 GB,不能放下完整模型,于是作者将模型某些 Layer 劈成两份,分别在两个 GPU 上计算)。
做个简单算术题。仍像上篇文章一样假设卷积层输入 feature map 尺寸为
如果平均分为 g 个组进行分组卷积,如下图(a)所示,输入 feature map 、卷积核、输出 feature map 平分为 3 个 group,在每个 group 内部仍采用经典卷积层计算,group 与 group 之间没有交叠。则:
每个组的输入 feature map 尺寸变为
/g 输出 feature map 尺寸变为
/g 卷积核尺寸变为
/g^2 每个组卷积计算量为:
/g^2 所有组卷积计算量累计为:
/g 所有组参数量累计为:
/g
结论:分组卷积(g>1)计算量和参数量相比未分组卷积(g=1)降低为 1/g。
当 g=M=N 时,刚好等价于深度方向卷积(depthwise convolution)。
分组卷积和通道重排示意图
2.2 通道重排
如上图(a),如果多个分组卷积堆叠起来,每个输出通道只能从有限输入通道获得信息,即一个组的输出只和这个组的输入有关,限制了模型表达能力(贸易保护主义不利于世界经济也是同样的道理)。
为了打破组间隔阂,(b) 提供了一种解决方案,每个 group 输入都有来自其他 group 上一层的输出。(c) 是一种较高效的实现,将每个 group 分为更小的 subgroup,通过 channel shuffle 操作,将每个 group 输出分散到每个 group 下一层输入。假设 GConv1 有 g x n 个输出通道,则首先将输出通道变维为 (g, n),再转置,最后展开为一维送入下一个层,这样就实现了通道重排。
通道重排操作是可微分的,能嵌入到网络中进行误差反向传播实现端到端训练。通道重排是 ShuffleNet 的精髓。
2.3 基本单元
基于上面的分组卷积和通道重排设计的 ShuffleNet 基本单元如下图所示:
该单元结构上采用残差网络(ResNet)【3】并做出了一系列改进来提升模型的效率:
(a) 使用深度方向卷积替换原有的 3x3 卷积;
(b) 将原先结构中前后两个 1x1 逐点卷积变为分组卷积,第一个分组卷积之后添加通道重排操作,降低卷积运算的跨通道计算量;
(c) 带下采样功能的 ShuffleNet;
经过以上调整,可以发现 ShuffleNet 在 MobileNet 基础上进一步降低了 1x1 卷积计算量与参数量(1/g),而据统计 1x1 卷积占据了整个 MobileNet 网络计算量和参数量的绝大部分(>90%、>70%),因此 ShuffleNet 这一改进是很有价值的。
2.4 完整网络
利用 ShuffleNet 单元构建的完整 ShuffeNet 网络模型如下表所示。
完整网络模型由16 个 ShuffleNet 基本单元堆叠而成,分三个阶段(Stage2 ~ Stage4),每阶段特征图空间尺寸减半,而通道数翻倍。整个模型的总计算量约为 140 MFLOPs。
通过 g=1, 2, 3, 4, 8 可以发现分组数值越大,保持总计算量不变时允许的通道数越多,网络越“宽”,潜在地有利于网络编码更多的信息,提升模型的识别能力。
通过增加一个网络“宽度”缩放因子 s 超参数,控制各层通道数目,可实现更精简的模型。
原网络 ShuffleNet 1x 计算复杂度为 140 MFLOPs
0.5x 将宽度缩减一半,计算量下降到 38 MFLOPs;
0.25x 将宽度缩减到 1/4,计算量下降到 13 MFLOPs;
3. 结果
上述模型分类结果如下表所示:
先纵向对比,发现模型越小,分类错误率越高,这符合直觉;
再横向对比,发现 g > 1 分类错误率低于 g = 1,即分组卷积可以提高模型分类准确率,但 g 也不是越大越好,对小模型而言,增大宽度可以提高模型分类准确率;
接着看通道重排对分类性能的影响,如下表所示:
对比发现使用通道重排相比不使用有1到4个百分点提升,分组数 g = 8 情况下尤其明显。
最后对比下 MobileNet 和 ShuffleNet,如下表所示。
发现在各个不同计算复杂度情况下,ShuffleNet 分类错误率都要低于 MobileNet,小模型上表现更为突出(低 4~6 个百分点),充分证明了 ShuffleNet 的强大和有效。
同几个流行的网络模型相比,在保持分类错误率相差无几情况下,大幅降低计算量(AlexNet:1/18,SqeezeNet:1/21,VGG-16:1/30)。
上表是在 ARM 处理器上的速度实测结果,可以发现 ShuffleNet 0.5x 模型与 AlexNet 分类错误率差不多,但实测速度有 13 倍提升(理论上有 18 倍),进一步说明 ShuffleNet 适合在移动设备上部署。
4. 实现
目前 Github 上 ShuffleNet 有各种实现,大部分基于 Caffe 框架,也有基于 PyTorch。TensorFlow 官方目前还没有这方面支持【4】。
使用 TensorFlow 提供的 tf.slice()、tf.concat() 可以实现分组操作,tf.reshape() 和 tf.transpose() 函数可以实现通道重排。感兴趣的读者可以试着实现 TensorFlow 版的 ShuffleNet。
5. 参考文献
【1】ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices, https://arxiv.org/abs/1707.01083
【2】Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In Advances in neural information processing systems, pages 1097–1105, 2012.
【3】 Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European Conference on Computer Vision, pages 630–645. Springer, 2016.
【4】https://github.com/tensorflow/models/issues/2035
如果你觉得本文对你有帮助,请关注公众号,将来会有更多更好的文章推送!