shufflenet v1
下图a是一个基于depthwise的残差网络, shufflenet基于这个基础模块修改成了b和c.
b为stride=1的block, c为stride=2的block. 与基础模块a相比, 使用GConv(g=3)替代了原始conv, 并且插入了channel shuffle模块. 对于stride=2的c block, skip connection分支使用avg conv来downscale.
- channel shuffle部分
def shuffle_channels(x, groups):
"""shuffle channels of a 4-D Tensor"""
batch_size, channels, height, width = x.size()
assert channels % groups == 0
channels_per_group = channels // groups
# split into groups
x = x.view(batch_size, groups, channels_per_group,
height, width)
# transpose 1, 2 axis
x = x.transpose(1, 2).contiguous()
# reshape into orignal
x = x.view(batch_size, channels, height, width)
return x
block c的实现:
class ShuffleNetUnitB(nn.Module):
"""ShuffleNet unit for stride=2"""
def __init__(self, in_channels, out_channels, groups=3):
super(ShuffleNetUnitB, self).__init__()
out_channels -= in_channels
assert out_channels % 4 == 0
bottleneck_channels = out_channels // 4
self.groups = groups
self.group_conv1 = nn.Conv2d(in_channels, bottleneck_channels,
1, groups=groups, stride=1)
self.bn2 = nn.BatchNorm2d(bottleneck_channels)
self.depthwise_conv3 = nn.Conv2d(bottleneck_channels,
bottleneck_channels,
3, padding=1, stride=2,
groups=bottleneck_channels)
self.bn4 = nn.BatchNorm2d(bottleneck_channels)
self.group_conv5 = nn.Conv2d(bottleneck_channels, out_channels,
1, stride=1, groups=groups)
self.bn6 = nn.BatchNorm2d(out_channels)
def forward(self, x):
out = self.group_conv1(x)
out = F.relu(self.bn2(out))
out = shuffle_channels(out, groups=self.groups)
out = self.depthwise_conv3(out)
out = self.bn4(out)
out = self.group_conv5(out)
out = self.bn6(out)
x = F.avg_pool2d(x, 3, stride=2, padding=1)
out = F.relu(torch.cat([x, out], dim=1))
return out
shufflenet v2
下图左侧ab为v1, 右侧为v2, 作者认为看一个网络的耗时不仅要看FLOPs, 还要看内存读取和gpu执行效率, 并提出了一些高效网络的设计准则.
G1). 使用输入通道和输出通道相同的卷积操作( 此时MAC(内存访问代价)最小);
G2). 谨慎使用分组卷积(MAC与分组数量g成正比, 所以应该尽量使用小的分组);
G3). 减少网络分支数(分支越多, 分支之间的卷积加载和分支同步就越耗时);
G4). 减少element-wise操作(element wise操作更耗时, 所以shufflenetv2使用了concat)。
注意上面提到的MAC为内存访问代价, 而不是计算量.
对比看来, v2将channel shuffle 放在了block最后位置, 两个分支的融合使用concat而不是 elementwise的add.
参考资料:
https://zhuanlan.zhihu.com/p/51566209
CNN模型之ShuffleNet v1