【他山之石】PyTorch实现ShuffleNet-v2亲身实践

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。

作者：Cai Yichao

地址：https://www.zhihu.com/people/YichaoCai

代码已同步到Github：https://github.com/EasonCai-Dev/torch_backbones

01

论文关键信息

ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design

论文链接：https://arxiv.org/abs/1807.11164

论文主要提出了ShuffleNet-v2的轻量级网络结构，并针对如今CNN网络常用的深度分离卷积（depthwise separable convolutions）、分组卷积（group convolutions）等进行了讨论，在网络结构上ShuffleNet的结构设计上看到了很多之前SoAT网络的启发。

1.1 设计高效的网络

除了提出网络结构之外，文章对设计高效的网络提供了指导建议：

G1) Equal channel width minimizes memory access cost (MAC). 在网络设计时尽量使输入输出的通道数相同以减少MAC，降低计算瓶颈并提高推理速度。

这一点，现在的SoAT网络结构确实已经可以看到这种思想的趋势，比如CSP-ResNe(X)t，CSP-DarkNet中取消bottleneck的操作, etc.

G2) Excessive group convolution increases MAC. 由于能够减少FLOPs，分组卷积自ResNeXt以来成为了很多网络设计的核心单元。但是过多地使用分组卷积会增加MAC。

设计网络时需要考虑使用场景，若需要较小的MAC以及需要从网络的层面提升推理速度时，应该要权衡分组卷积的使用频率。

G3) Network fragmentation reduces degree of parallelism. 网络过度地“碎片化”会降低计算的并行度。

比如Inception系列中的block（尤其是V3之前的版本）重复度不高，并且block内部结构不相同，这相比于ResNe(X)t系列来说，除了显得不那么优雅之外，会增加很多碎片化地操作，而不是大量的相同操作，这样使得计算并行度降低。这样，随着网络加深，即使网络能够很好地避免梯度问题，但其推理速度会下降。

G4) Element-wise operations are non-negligible. Element-wise操作，比如非线性激活函数、张量相加、添加偏置以及depthwise conv等，相比于卷积操作来说其计算量小很多，但是也是不能忽略的，在设计网络时应该要考虑这些因素。

1.2 网络结构

这篇文章的主题是ShuffleNet-v2网络结构的PyTorch实现，如果需要从论文中获取更多的性能对比、思路启发，建议还是阅读以下原文。那么我们直接切入正题。

1.2.1 Channel Shuffle

Channel Shuffle是ShuffleNet-v1（https://arxiv.org/abs/1707.01083）论文的创新点，可以再保持网络精度的情况下提升计算速度，其思想是增加通道间的信息流动，示意图如下。如图，在输入到分组卷积之前，将特征按照分组卷积的group数量进行分组，并且将组间的特征进行打乱，这样就提高了特征的信息重用。

1.2.1 ShuffleNet-v2

下图是论文中给出的ShuffleNet-v2的block结构，(c)为basic unit， (d)为down-sampling unit。

1. Down-sampling unit: 由于Concat的存在，输出通道为输入通道的两倍，并且使用3x3的DWConv进行降维。

这个结构在设计上也遵循了：使用卷积进行降维时，通道进行翻倍来维持特征量的思想。忘了在哪篇论文中看到过，我之前的博客应该有介绍这个思想。

2. Basic unit：在这个结构中可以看到很多网络的影子，分组卷积，Concat，CSPD等都有。n.b. 在block中Channel Shuffle被放在了Concat之后，而不是像v1一样放在分组卷积前。这点可以这样理解，两个basic unit进行串接，下一个unit处理前会将特征进行进行通道分离操作，分成两组，所以在Channel Shuffle的时候将分成两组进行Shuffle是讲得通的。除了上图的基础版本之外，作者还结合Residual和SE结构（https://zhuanlan.zhihu.com/p/263537813）整合了下图3种block。

其中的DWConv表示的是Depthwise convolution，其输入和输入通道数相同，并且每个卷积核只与一个特征通道进行卷积，其使用线性激活函数。 另外，惯常的是当使用Residual block时，在shortcut之前也是采用线性激活函数而非ReLU。

3. 网络架构

网络整体架构如下表。论文给出了具有四种不同比例的结构，其不同主要体现在Stage2的输出通道数量上面。网络整体上每个Stage内部由一个Down-sampling unit和数个Basic unit重复串接，除了Stage2之外，每个basic unit的输入输出通道相同，每个stage的输入输出特征量相同，但是其通道数量翻倍，特征尺寸减半。

02

PyTorch实现

2.1 Channel Shuffle

通过Reshape转置后恢复原来的张量尺寸，来取得Channel Shuffle的效果，代码如下：

def shuffle_chnls(x, groups=2):
    """Channel Shuffle"""

    bs, chnls, h, w = x.data.size()
    if chnls % groups:
        return x
    chnls_per_group = chnls // groups
    x = x.view(bs, groups, chnls_per_group, h, w)
    x = torch.transpose(x, 1, 2).contiguous()
    x = x.view(bs, -1, h, w)
    return x

2.2 Block结构

2.2.1 BN_Conv

把之前写的BN_Conv代码改造了一下，以应对有些时候采用线性激活函数的情况。默认采用ReLU，可以通过传参控制改变激活函数类型，输入None表示线性激活，代码如下：

class BN_Conv2d(nn.Module):
    """
    BN_CONV, default activation is ReLU
    """

    def __init__(self, in_channels: object, out_channels: object, kernel_size: object, stride: object, padding: object,
                 dilation=1, groups=1, bias=False, activation=nn.ReLU(inplace=True)) -> object:
        super(BN_Conv2d, self).__init__()
        layers = [nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride,
                            padding=padding, dilation=dilation, groups=groups, bias=bias),
                  nn.BatchNorm2d(out_channels)]
        if activation is not None:
            layers.append(activation)
        self.seq = nn.Sequential(*layers)

    def forward(self, x):
        return self.seq(x)

2.2.2 Down-sampling Unit

DWConv实现时将groups属性设置为和out_channels相同即可。代码如下：

class DSampling(nn.Module):
    """Spatial down sampling of SuffleNet-v2"""

    def __init__(self, in_chnls, groups=2):
        super(DSampling, self).__init__()
        self.groups = groups
        self.dwconv_l1 = BN_Conv2d(in_chnls, in_chnls, 3, 2, 1,  # down-sampling, depth-wise conv.
                                   groups=in_chnls, activation=None)
        self.conv_l2 = BN_Conv2d(in_chnls, in_chnls, 1, 1, 0)
        self.conv_r1 = BN_Conv2d(in_chnls, in_chnls, 1, 1, 0)
        self.dwconv_r2 = BN_Conv2d(in_chnls, in_chnls, 3, 2, 1, groups=in_chnls, activation=None)
        self.conv_r3 = BN_Conv2d(in_chnls, in_chnls, 1, 1, 0)

    def forward(self, x):
        # left path
        out_l = self.dwconv_l1(x)
        out_l = self.conv_l2(out_l)

        # right path
        out_r = self.conv_r1(x)
        out_r = self.dwconv_r2(out_r)
        out_r = self.conv_r3(out_r)

        # concatenate
        out = torch.cat((out_l, out_r), 1)
        out = shuffle_chnls(out, self.groups)
        return shuffle_chnls(out, self.groups)

2.2.3 Basic Unit

代码如下：

class BasicUnit(nn.Module):
    """Basic Unit of ShuffleNet-v2"""

    def __init__(self, in_chnls, out_chnls, is_se=False, is_residual=False, c_ratio=0.5, groups=2):
        super(BasicUnit, self).__init__()
        self.is_se, self.is_res = is_se, is_residual
        self.l_chnls = int(in_chnls * c_ratio)
        self.r_chnls = in_chnls - self.l_chnls
        self.ro_chnls = out_chnls - self.l_chnls
        self.groups = groups

        # layers
        self.conv1 = BN_Conv2d(self.r_chnls, self.ro_chnls, 1, 1, 0)
        self.dwconv2 = BN_Conv2d(self.ro_chnls, self.ro_chnls, 3, 1, 1,  # same padding, depthwise conv
                                 groups=self.ro_chnls, activation=None)
        act = None if self.is_res else nn.ReLU(inplace=True)
        self.conv3 = BN_Conv2d(self.ro_chnls, self.ro_chnls, 1, 1, 0, activation=act)
        if self.is_se:
            self.se = SE(self.ro_chnls, 16)
        if self.is_res:
            self.shortcut = nn.Sequential()
            if self.r_chnls != self.ro_chnls:
                self.shortcut = BN_Conv2d(self.r_chnls, self.ro_chnls, 1, 1, 0, activation=None)

    def forward(self, x):
        x_l = x[:, :self.l_chnls, :, :]
        x_r = x[:, self.l_chnls:, :, :]

        # right path
        out_r = self.conv1(x_r)
        out_r = self.dwconv2(out_r)
        out_r = self.conv3(out_r)
        if self.is_se:
            coefficient = self.se(out_r)
            out_r *= coefficient
        if self.is_res:
            out_r += self.shortcut(x_r)

        # concatenate
        out = torch.cat((x_l, out_r), 1)
        return shuffle_chnls(out, self.groups)

2.3 网络整体架构和API

通过参数控制是否使用SE和Residual结构，以及控制网络的尺寸，在{0.5, 1, 1.5, 2}四个尺寸中进行选择。

代码如下：

class ShuffleNet_v2(nn.Module):
    """ShuffleNet-v2"""

    _defaults = {
        "sets": {0.5, 1, 1.5, 2},
        "units": [3, 7, 3],
        "chnl_sets": {0.5: [24, 48, 96, 192, 1024],
                      1: [24, 116, 232, 464, 1024],
                      1.5: [24, 176, 352, 704, 1024],
                      2: [24, 244, 488, 976, 2048]}
    }

    def __init__(self, scale, num_cls, is_se=False, is_res=False) -> object:
        super(ShuffleNet_v2, self).__init__()
        self.__dict__.update(self._defaults)
        assert (scale in self.sets)
        self.is_se = is_se
        self.is_res = is_res
        self.chnls = self.chnl_sets[scale]

        # make layers
        self.conv1 = BN_Conv2d(3, self.chnls[0], 3, 2, 1)
        self.maxpool = nn.MaxPool2d(3, 2, 1)
        self.stage2 = self.__make_stage(self.chnls[0], self.chnls[1], self.units[0])
        self.stage3 = self.__make_stage(self.chnls[1], self.chnls[2], self.units[1])
        self.stage4 = self.__make_stage(self.chnls[2], self.chnls[3], self.units[2])
        self.conv5 = BN_Conv2d(self.chnls[3], self.chnls[4], 1, 1, 0)
        self.globalpool = nn.AdaptiveAvgPool2d((1, 1))
        self.body = self.__make_body()
        self.fc = nn.Linear(self.chnls[4], num_cls)

    def __make_stage(self, in_chnls, out_chnls, units):
        layers = [DSampling(in_chnls),
                  BasicUnit(2 * in_chnls, out_chnls, self.is_se, self.is_res)]
        for _ in range(units-1):
            layers.append(BasicUnit(out_chnls, out_chnls, self.is_se, self.is_res))
        return nn.Sequential(*layers)

    def __make_body(self):
        return nn.Sequential(
            self.conv1, self.maxpool, self.stage2, self.stage3,
            self.stage4, self.conv5, self.globalpool
        )

    def forward(self, x):
        out = self.body(x)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return F.softmax(out)


"""
API
"""


def shufflenet_0_5x(num_classes=1000):
    return ShuffleNet_v2(0.5, num_classes)


def shufflenet_0_5x_se(num_classes=1000):
    return ShuffleNet_v2(0.5, num_classes, is_se=True)


def shufflenet_0_5x_res(num_classes=1000):
    return ShuffleNet_v2(0.5, num_classes, is_res=True)


def shufflenet_0_5x_se_res(num_classes=1000):
    return ShuffleNet_v2(0.5, num_classes, is_se=True, is_res=True)


def shufflenet_1x(num_classes=1000):
    return ShuffleNet_v2(1, num_classes)


def shufflenet_1x_se(num_classes=1000):
    return ShuffleNet_v2(1, num_classes, is_se=True)


def shufflenet_1x_res(num_classes=1000):
    return ShuffleNet_v2(1, num_classes, is_res=True)


def shufflenet_1x_se_res(num_classes=1000):
    return ShuffleNet_v2(1, num_classes, is_se=True, is_res=True)


def shufflenet_1_5x(num_classes=1000):
    return ShuffleNet_v2(1.5, num_classes)


def shufflenet_1_5x_se(num_classes=1000):
    return ShuffleNet_v2(1.5, num_classes, is_se=True)


def shufflenet_1_5x_res(num_classes=1000):
    return ShuffleNet_v2(1.5, num_classes, is_res=True)


def shufflenet_1_5x_se_res(num_classes=1000):
    return ShuffleNet_v2(1.5, num_classes, is_se=True, is_res=True)


def shufflenet_2x(num_classes=1000):
    return ShuffleNet_v2(2, num_classes)


def shufflenet_2x_se(num_classes=1000):
    return ShuffleNet_v2(2, num_classes, is_se=True)


def shufflenet_2x_res(num_classes=1000):
    return ShuffleNet_v2(2, num_classes, is_res=True)


def shufflenet_2x_se_res(num_classes=1000):
    return ShuffleNet_v2(2, num_classes, is_se=True, is_res=True)

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

直播预告

历史文章推荐

分享、点赞、在看，给个三连击呗！