介绍
在论文 ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices 中,作者提出了 Group convolution 帮助降低计算复杂度;但是使用 Group convolution 会有副作用,故在此基础上,论文提出 Channel shuffle 帮助信息流通。基于这两种技术,论文构建了一个名为 ShuffleNet 的高效架构,相比于其他先进模型,对于给定的计算复杂度预算,ShuffleNet 允许使用更多的特征映射通道,在小型网络上有助于编码更多信息。
Group convolution
Group convolution 的实质就是将标准卷积过程分为 g g g 个独立的组,分别计算,最后再合并起来。即:
- 假设输入的形状为 H × W × C H\times W \times C H×W×C,用 k k k 个 h × w h\times w h×w 的卷积核对其进行卷积操作;
- 把输入分为 g g g 组,每组形状为 H × W × ( C / g ) H\times W \times (C/g) H×W×(C/g)(假设可整除);
- 把卷积核也分为 g g g 组,每组为 k / g k/g k/g (假设可整除)个 h × w h\times w h×w 卷积核;
- 按顺序,每组的输入和该组内的卷积核分别做标准卷积操作,输出 g g g 组形状为 H ′ × W ′ × ( k / g ) H'\times W'\times (k/g) H′×W′×(k/g);
- 将这 g g g 组特征合并起来,得到最终形状为 H ′ × W ′ × k H'\times W'\times k H′×W′×k 的特征。
下图(来源于白裳)是分两组时的举例:
使用 Group convolution 的好处是大大减少了网络的参数量:
对于标准卷积,假设输入的形状为
H
×
W
×
C
H\times W \times C
H×W×C,用
k
k
k 个
h
×
w
h\times w
h×w 的卷积核对其进行卷积操作,其参数量为:
h
×
w
×
C
×
k
h\times w\times C\times k
h×w×C×k
对于 Group convolution,输入和卷积核不变,其参数量为:
h
×
w
×
C
/
g
×
k
/
g
×
g
=
h
×
w
×
C
×
k
/
g
h\times w\times C/g\times k/g \times g=h\times w\times C\times k/g
h×w×C/g×k/g×g=h×w×C×k/g
这比原来缩小了
g
g
g 倍。
Channel shuffle
因为 Group convolution 将通道分开进行卷积操作,所以会导致通道具有局部性,这样对模型的泛化能力有所损失,影响模型的准确率。因此我们使用 Channel shuffle 对 Group convolution 的卷积结果在通道维度进行打乱,即将输出的通道进行重新分组。具体的实现如下图(c)所示。
其具体步骤为:
- 假设输入的形状为 h × w × c h\times w\times c h×w×c;
- 假设将输入层分成 g g g 组,则每组的通道数是 c / g c/g c/g;
- 将通道那个维度拆分为两个维度,即将形状变为 h × w × g × c / g h\times w\times g\times c/g h×w×g×c/g;
- 对最后那两个拆分出来的维度进行转置,得到形状为 h × w × c / g × g h\times w\times c/g\times g h×w×c/g×g;
- 对最后两个维度重新 reshape 成一个维度,重新得到形状为 h × w × c h\times w\times c h×w×c。
ShuffleNet Unit
下面三幅图中:
-
图 a) 是标准的 bottleneck unit,和 MobileNet V2 中的一样,其中,第二个 1x1 Conv 的目的是把输出的通道数恢复到和输入一样,这样才能进行相加操作,即 shortcut path。
-
图 b) 是一个 ShuffleNet Unit,它与图 a) 的区别有:
- 1、把图 a) 中前后两个 1x1 Conv 采用了 Group convolution 的形式,称之为 Pointwise group convolution;
- 2、Channel shuffle 被应用在第一个 Pointwise group convolution 之后。
-
图 c) 也是一个 ShuffleNet Unit,它跟图 b) 的区别是:
- 1、中间一个 3x3 的 Depthwise convolution 的步长是 2 而不是 1;
- 2、在支路上加一个步长为 2 的 3x3 的平均池化层;
- 3、把相加操作变成合并操作,这样可以在计算成本增加不多的情况下增多通道数。
ShuffleNet 总体结构
在三个 stage (stage 2, 3,4) 中,都是先使用步长为 2 的 ShuffleNet Unit,又堆叠若干个个步长为 1 的 ShuffleNet Unit。
另外,某一个 stage 的输出通道数是上一个 stage 的输出通道数的两倍,且在每个 Unit 里,设置 Bottleneck 的通道数是输出通道数的四分之一。
【注】在 ShuffleNet 里,Depthwise convolution 只用在 Bottleneck 里,这是因为虽然 Depthwise convolution 在理论上并不复杂,但它很难有效地被应用到低功率移动设备上。
除此之外,我们使用 s 来表示某种 ShuffleNet 中卷积核的个数,如下图所示:
ShuffleNet sx 都是基于 ShuffleNet 1x 来说的,表示卷积核的个数(通道数)是 ShuffleNet 1× 中的 s 倍。
代码实现
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 group_conv(inputs, filters, kernel, strides, num_groups):
conv_side_layers_tmp = tf.split(inputs, num_groups ,axis=3)
conv_side_layers = []
for layer in conv_side_layers_tmp:
conv_side_layers.append(tf.keras.layers.Conv2D(filters//num_groups, kernel, strides, padding='same')(layer))
x = tf.concat(conv_side_layers, axis=-1)
return x
def conv(inputs, filters, kernel_size, stride=1, activation=False):
x = tf.keras.layers.Conv2D(filters, kernel_size, stride, padding='same')(inputs)
x = tf.keras.layers.BatchNormalization()(x)
if activation:
x = tf.keras.layers.Activation('relu')(x)
return x
def depthwise_conv_bn(inputs, kernel_size, stride=1):
x = tf.keras.layers.DepthwiseConv2D(kernel_size=kernel_size,
strides=stride,
padding='same')(inputs)
x = tf.keras.layers.BatchNormalization()(x)
return x
def ShuffleNetUnitA(inputs, num_groups):
in_channels = inputs.shape[-1]
out_channels = in_channels
bottleneck_channels = out_channels // 4
x = group_conv(inputs, bottleneck_channels, kernel=1, strides=1, num_groups=num_groups)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Activation('relu')(x)
x = channel_shuffle(x, num_groups)
x = depthwise_conv_bn(x, kernel_size=3, stride=1)
x = group_conv(x, out_channels, kernel=1, strides=1, num_groups=num_groups)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.add([inputs, x])
x = tf.keras.layers.Activation('relu')(x)
return x
def ShuffleNetUnitB(inputs, out_channels, num_groups):
in_channels = inputs.shape[-1]
out_channels -= in_channels
bottleneck_channels = out_channels // 4
x = group_conv(inputs, bottleneck_channels, kernel=1, strides=1, num_groups=num_groups)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Activation('relu')(x)
x = channel_shuffle(x, num_groups)
x = depthwise_conv_bn(x, kernel_size=3, stride=2)
x = group_conv(x, out_channels, kernel=1, strides=1, num_groups=num_groups)
x = tf.keras.layers.BatchNormalization()(x)
y = tf.keras.layers.AvgPool2D(pool_size=3, strides=2, padding='same')(inputs)
x = tf.concat([y, x], axis=-1)
x = tf.keras.layers.Activation('relu')(x)
return x
def stage(inputs, out_channels, num_groups, n):
x = ShuffleNetUnitB(inputs, out_channels, num_groups)
for _ in range(n):
x = ShuffleNetUnitA(x, num_groups)
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, num_groups, n=3)
x = stage(x, first_stage_channels*2, num_groups, n=7)
x = stage(x, first_stage_channels*4, num_groups, n=3)
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])