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_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} MAC≥2hwB+hwB
仅当 c 1 = c 2 c_1=c_2 c1=c2 时,MAC 取最小值,这个理论分析也通过实验得到证实,如下表所示,通道比为 1:1 时速度更快。
- (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 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 在全局池化之前增加了一个卷积层。
分类效果
代码实现
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])