ShuffleNet V2 网络结构的原理与 Tensorflow2.0 实现

ShuffleNet V2 的设计准则

目前衡量模型复杂度的一个通用指标是 FLOPs，具体指的是计算在网络中的乘法操作和加法操作的数量，但是这却是一个间接指标，因为它不完全等同于速度。如下图中的（c）和（d），可以看到具有相同 FLOPs 的两个模型，其速度在不同系统架构下却存在差异。这种不一致主要归结为两个原因，首先除了 FLOPs 影响速度之外，还有内存使用量（memory access cost, MAC）等，这不能忽略，且对于 GPUs 来说可能会是瓶颈。另外模型的并行程度也影响速度，并行度高的模型速度相对更快。另外一个原因，模型在不同平台上的运行速度是有差异的，如 GPU 和 ARM，而且采用不同的库也会有影响。

据此，作者在特定的平台下研究 ShuffleNet V1 和 MobileNet V2 的运行时间，并结合理论与实验得到了 4 条实用的指导原则：

• （G1）同等通道大小最小化 MAC：对于轻量级 CNN 网络，常采用 Depthwise separable convolutions，其中 Pointwise convolution 即 1x1 Conv 的复杂度最大。这里假定输入和输出特征的通道数分别为 $c_1$ 和 $c_2$ ，特征图的空间大小为 $h\times w$，那么 1x1 Conv 的 FLOPs 为 $B=hw{c}_1{c}_2$。对应的 MAC 为 $hw(c_1+c_2)+c_1{c}_2$（$hw{c}_1$ 是输入的内存大小，$hw{c}_2$ 是输出的内存大小，$c_1{c}_2$ 是 $c_2$ 个 1x1x$c_1$ 卷积核中的权重所占内存大小），根据均值不等式，固定 $B$ 时，MAC 存在下限（令 $c_2=\frac{B}{hw{c}_1}$）：
  $$MAC\ge 2\sqrt{hwB}+\frac{B}{hw}$$
  仅当 $c_1=c_2$ 时，MAC 取最小值，这个理论分析也通过实验得到证实，如下表所示，通道比为 1:1 时速度更快。
  
• （G2）过量使用 Group convolution 会增加 MAC：Group convolution 是常用的设计组件，因为它可以减少复杂度却不损失模型容量。但是这里发现，分组过多会增加 MAC。对于 Group convolution，FLOPs 为 $B=hw{c}_1{c}_2/g$ （其中 $g$ 是分出来的组数），而对应的 MAC 为 $hw(c_1+c_2)+c_1{c}_2/g$。如果固定输入 $h\times w\times c_1$ 以及 $B$，那么 MAC 为
  $$MAC=hw{c}_1+Bg/c_1+B/hw$$可以看到，当 $g$ 增加时，MAC 会同时增加。这点也通过实验证实，所以明智之举是不要使用太大 $g$ 的 Group convolution。
• （G3）网络碎片化会降低并行度：一些网络如 Inception，以及 Auto ML 自动产生的网络 NASNET-A，它们倾向于采用“多路”结构，即将几个卷积层和池化层放到一个 Block 中，这很容易造成网络碎片化，从而减低模型的并行度，使得速度变慢。
• （G4）不能忽略元素级操作：对于 ReLU 和 Add 操作，虽然它们的 FLOPs 较小，但是却需要较大的 MAC。这里实验发现如果将 ResNet 中残差单元中的 ReLU 和 shortcut 移除的话，速度有 20% 的提升。

ShuffleNet V2 的改进

根据前面的 4 条准则，作者分析了 ShuffleNet V1 设计的不足，并在此基础上改进得到了 ShuffleNet V2，两者模块上的对比如下图（a 和 b 是 ShuffleNet V1 中的两种 units，c 和 d 是 ShuffleNet V2 中的两种 units）所示：

ShuffleNet V1 中有以下缺陷：

• 采用了类似 ResNet 中的Bottleneck layer，使得输入和输出通道数不同，违背了 G1 原则；
• 大量使用 1x1 Group convolution，违背了 G2 原则，同时使用过多的组，也违背了 G3 原则；
• 短路连接中存在大量的 Add 运算，违背了 G4 原则。

为了改善 V1 中的缺陷，V2 版本引入了一种新的运算：Channel split。具体来说，在开始时先将通道数为 $c$ 的输入特征图在通道维度分成两个分支，每个分支的通道数都是 $c/2$。这样做的好处是：

• 左边分支做同等映射，右边的分支包含 3 个连续的卷积，并且输入和输出通道相同，这符合 G1；
• 两个 1x1 Conv 不再是 Group convolution，这符合 G2；
• 两个分支的输出不再用 Add 相加，而是 Concat 在一起，这符合 G4。

【注】对于下采样模块，即 ShuffleNet V2 中的第二个 unit，不再有 Channel split，而是每个分支都是直接复制一份输入，每个分支都有步长为 2 的下采样，最后 Concat 在一起后，特征图空间大小减半，但是通道数翻倍。

网络结构

ShuffleNet V2 的整体结构如下表所示，基本与 V1 类似：

【注】ShuffleNet V2 在全局池化之前增加了一个卷积层。

分类效果

代码实现

import tensorflow as tf

def channel_shuffle(inputs, num_groups):
    
    n, h, w, c = inputs.shape
    x_reshaped = tf.reshape(inputs, [-1, h, w, num_groups, c // num_groups])
    x_transposed = tf.transpose(x_reshaped, [0, 1, 2, 4, 3])
    output = tf.reshape(x_transposed, [-1, h, w, c])
    
    return output

def conv(inputs, filters, kernel_size, strides=1):
    
    x = tf.keras.layers.Conv2D(filters, kernel_size, strides, padding='same')(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation('relu')(x)
        
    return x

def depthwise_conv_bn(inputs, kernel_size, strides=1):
    
    x = tf.keras.layers.DepthwiseConv2D(kernel_size=kernel_size, 
                                        strides=strides, 
                                        padding='same')(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
        
    return x

def ShuffleNetUnitA(inputs, out_channels):
    
    shortcut, x = tf.split(inputs, 2, axis=-1)
    
    x = conv(inputs, out_channels // 2, kernel_size=1, strides=1)
    x = depthwise_conv_bn(x, kernel_size=3, strides=1)
    x = conv(x, out_channels // 2, kernel_size=1, strides=1)
    
    x = tf.concat([shortcut, x], axis=-1)
    x = channel_shuffle(x, 2)
    
    return x

def ShuffleNetUnitB(inputs, out_channels):
    
    shortcut = inputs
    
    in_channels = inputs.shape[-1]
    
    x = conv(inputs, out_channels // 2, kernel_size=1, strides=1)
    x = depthwise_conv_bn(x, kernel_size=3, strides=2)
    x = conv(x, out_channels-in_channels, kernel_size=1, strides=1)
    
    shortcut = depthwise_conv_bn(shortcut, kernel_size=3, strides=2)
    shortcut = conv(shortcut, in_channels, kernel_size=1, strides=1)
    
    output = tf.concat([shortcut, x], axis=-1)
    output = channel_shuffle(output, 2)
    
    return output

def stage(inputs, out_channels, n):
    
    x = ShuffleNetUnitB(inputs, out_channels)
    
    for _ in range(n):
        x = ShuffleNetUnitA(x, out_channels)
        
    return x

def ShuffleNet(inputs, first_stage_channels, num_groups):
    x = tf.keras.layers.Conv2D(filters=24, 
                               kernel_size=3, 
                               strides=2, 
                               padding='same')(inputs)
    x = tf.keras.layers.MaxPooling2D(pool_size=3, strides=2, padding='same')(x)
    
    x = stage(x, first_stage_channels, n=3)
    x = stage(x, first_stage_channels*2, n=7)
    x = stage(x, first_stage_channels*4, n=3)
    
    x = tf.keras.layers.Conv2D(filters=1024, 
                               kernel_size=1, 
                               strides=1, 
                               padding='same')(x)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dense(1000)(x)
    
    return x

inputs = np.zeros((1, 224, 224, 3), np.float32)
ShuffleNet(inputs, 144, 1).shape

TensorShape([1, 1000])

参考资料

ShuffleNetV2：轻量级CNN网络中的桂冠