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

文章目录

  • ShuffleNet V2 的设计准则
  • ShuffleNet V2 的改进
  • 网络结构
  • 分类效果
  • 代码实现
  • 参考资料

ShuffleNet V2 的设计准则

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

ShuffleNet V2 网络结构的原理与 Tensorflow2.0 实现_第1张图片
据此,作者在特定的平台下研究 ShuffleNet V1 和 MobileNet V2 的运行时间,并结合理论与实验得到了 4 条实用的指导原则:

  • (G1)同等通道大小最小化 MAC:对于轻量级 CNN 网络,常采用 Depthwise separable convolutions,其中 Pointwise convolution 即 1x1 Conv 的复杂度最大。这里假定输入和输出特征的通道数分别为 c 1 c_1 c1 c 2 c_2 c2 ,特征图的空间大小为 h × w h\times w h×w,那么 1x1 Conv 的 FLOPs 为 B = h w c 1 c 2 B=hwc_1c_2 B=hwc1c2。对应的 MAC 为 h w ( c 1 + c 2 ) + c 1 c 2 hw(c_1+c_2)+c_1c_2 hw(c1+c2)+c1c2 h w c 1 hwc_1 hwc1 是输入的内存大小, h w c 2 hwc_2 hwc2 是输出的内存大小, c 1 c 2 c_1c_2 c1c2 c 2 c_2 c2 个 1x1x c 1 c_1 c1 卷积核中的权重所占内存大小),根据均值不等式,固定 B B B 时,MAC 存在下限(令 c 2 = B h w c 1 c_2=\frac{B}{hwc_1} c2=hwc1B):
    M A C ≥ 2 h w B + B h w MAC\geq 2\sqrt{hwB}+\frac{B}{hw} MAC2hwB +hwB
    仅当 c 1 = c 2 c_1=c_2 c1=c2 时,MAC 取最小值,这个理论分析也通过实验得到证实,如下表所示,通道比为 1:1 时速度更快。
    ShuffleNet V2 网络结构的原理与 Tensorflow2.0 实现_第2张图片
  • (G2)过量使用 Group convolution 会增加 MAC:Group convolution 是常用的设计组件,因为它可以减少复杂度却不损失模型容量。但是这里发现,分组过多会增加 MAC。对于 Group convolution,FLOPs 为 B = h w c 1 c 2 / g B=hwc_1c_2/g B=hwc1c2/g (其中 g g g 是分出来的组数),而对应的 MAC 为 h w ( c 1 + c 2 ) + c 1 c 2 / g hw(c_1+c_2)+c_1c_2/g hw(c1+c2)+c1c2/g。如果固定输入 h × w × c 1 h\times w\times c_1 h×w×c1 以及 B B B,那么 MAC 为
    M A C = h w c 1 + B g / c 1 + B / h w MAC=hwc_1+Bg/c_1+B/hw MAC=hwc1+Bg/c1+B/hw可以看到,当 g g g 增加时,MAC 会同时增加。这点也通过实验证实,所以明智之举是不要使用太大 g g 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 V2 网络结构的原理与 Tensorflow2.0 实现_第3张图片
ShuffleNet V1 中有以下缺陷:

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

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

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

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

网络结构

ShuffleNet V2 的整体结构如下表所示,基本与 V1 类似:
ShuffleNet V2 网络结构的原理与 Tensorflow2.0 实现_第4张图片
【注】ShuffleNet V2 在全局池化之前增加了一个卷积层。

分类效果

ShuffleNet V2 网络结构的原理与 Tensorflow2.0 实现_第5张图片

代码实现

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网络中的桂冠

你可能感兴趣的:(Tensorflow,2.0,深度学习,网络,深度学习,卷积,卷积神经网络,tensorflow)