ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design

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

一、论文简介：

深度卷积神经网络(CNN)的架构经过多年的发展，变得更加准确和快速。除了精度，计算复杂度是另一个重要的考虑因素。现实中的任务通常旨在在有限的计算预算下获得最佳的精度。这推动了一系列朝着轻量级架构设计和更好的速度-精度折衷的工作，包括Xception、MobileNet、MobileNet V2、ShuffleNet和CondenseNet。分组卷积和深度可分离卷积是这些工作的关键。
为了度量计算复杂度，一个广泛使用的度量指标是浮点运算的数量（FLOPs）。然而，FLOPs是一个间接的指标。这是一个近似值，通常不等同于我们真正关心的直接指标，如速度或延迟。在之前的研究中，这种差异已被注意到。例如，MobileNet v2比NASNET-A快得多，但它们有基本相同的准确率。下图表明具有类似FLOPs的网络具有不同的速度。因此，使用FLOPs作为计算复杂度的唯一度量是不够的，可能会导致次优化设计。


更多比较见上图。ShuffleNet v2显然比其他三个网络快，尤其是在GPU上。例如，在500MFLOPs时，ShuffleNet v2比MobileNet v2快58%，比ShuffleNet v1快63%，比Xception快25%。在ARM上，ShuffleNet v1、Xception和ShuffleNet v2的速度相当;然而，MobileNet v2要慢得多，特别是在较小的FLOPs上。这是因为MobileNet v2拥有更高的MAC，这在移动设备上是显著的。
间接(FLOPs)和直接(速度)度量之间的差异可以归因于两个主要原因。首先，几个对速度有很大影响的重要因素并没有被FLOPs考虑在内。其中一个因素是内存访问成本(MAC)。在某些操作(如分组卷积)中，这种开销占了运行时的很大一部分。它可能成为计算能力强的设备（例如gpu）的瓶颈。在网络架构设计时，不应简单地忽略这一成本。另一个是平行度。在相同的FLOPs下，具有高并行度的模型可以比具有低并行度的模型快得多。其次，使用相同FLOPs的操作可能有不同的运行时间，这取决于平台。例如，张量分解在早期的研究中被广泛使用来加速矩阵乘法。尽管它减少了75%的FLOPs，但是在GPU上的分解速度甚至更慢。这是因为最新的CUDNN库专门针对3×3 conv进行了优化。我们不能肯定地认为3×3 conv比1×1 conv慢9倍。
根据这些观察结果，作者提出在有效的网络架构设计中应该考虑两个原则。首先，应该使用直接的指标(如速度)，而不是间接的指标(如FLOPs)。其次，该指标应在目标平台上进行评估。
在这项工作中，作者遵循这两个原则，并提出一个更有效的网络架构。在第二节中，作者首先分析了两个最先进的代表性网络的运行性能。然后，推导出四个指导原则，进行了有效的网络设计，其超越了仅仅考虑FLOPs的有效性。虽然这些指导原则是平台独立的，但在两个不同的平台(GPU和ARM)上进行了一系列受控实验，通过专门的代码优化来验证它们，确保结论是最先进的。
在第三节中，作者根据指南设计了一个新的网络结构。它的灵感来自于ShuffleNet，因此被称为ShuffleNet V2。通过第4节的验证实验，证明了该算法在两个平台上都比之前的网络更快、更准确。下图给出了比较的概述。例如，给定40M FLOPs的计算复杂度预算，ShuffleNet v2比ShuffleNet v1和MobileNet v2分别高3.5%和3.7%。

二、高效网络设计实用指南

为了开始研究，作者分析了两个最先进的网络，ShuffleNet v1和MobileNet v2的运行时性能。它们在ImageNet分类任务中既高效又准确。它们都广泛应用于手机等低端设备。虽然只分析了这两个网络，但作者注意到它们代表了当前的趋势。它们的核心是分组卷积和深度卷积，这也是其他最先进的网络的关键组件，如ResNeXt、Xception、MobileNet和CondenseNet。
模型整体运行时被分解为不同的操作，如下图所示。作者注意到FLOPs度量只考虑了卷积部分。虽然这部分占用的时间最多，但其他的操作，包括数据I/O、数据shuffle和元素级的操作(AddTensor、ReLU等)也占用了大量的时间。因此，FLOPs对实际运行时的估计不够精确。

基于这一观察，作者从几个不同的方面对运行时间(或速度)进行了详细的分析，并得出了几个有效的网络架构设计的实用指南。

（一）、通道宽度相等使内存访问成本(MAC)最小

现代网络通常采用深度可分离卷积，其中点卷积(即1×1卷积)占了大部分的复杂度。作者研究了1×1卷积的kernel shape。kernel shape由两个参数指定:输入通道数量$c_1$和输出通道数量$c_2$。设h和w为feature map的空间大小，则1×1卷积的FLOPs为B = $hw{c}_1{c}_2$。
为简单起见，作者假设计算设备中的缓存足够大，可以存储整个特征映射和参数。因此，内存访问成本(MAC)，即内存访问操作的次数，为MAC = $hw(c_1+c_2)+c_1{c}_2$。注意，公式的两部分分别对应于输入/输出特征映射和内核权重的内存访问。

从上式中我们可以看出，MAC有一个由FLOPs给出的下限。当输入和输出通道数相等时，它达到下界。
以上的结论是理论性的。实际上，许多设备上的缓存都不够大。此外，现代计算库通常采用复杂的阻塞策略来充分利用缓存机制。因此，实际MAC可能会与理论MAC有所偏差。为了验证上述结论，作者进行了如下实验。一个benchmark network是通过10个构件块重复堆叠而成的。每个块包含两个卷积层。第一个包含$c_1$输入通道和$c_2$输出通道，第二个则不用。
下表通过改变$c_1$: $c_2$比率报告了运行速度，同时确定了总的FLOPs。可以看出，当$c_1$: c_2$趋于1:1时，MAC变小，网络评估速度加快。

（二）、过多的分组卷积会增加MAC

分组卷积是现代网络体系结构的核心。它通过改变所有通道之间的密集卷积为稀疏(仅在通道组内)来降低计算复杂度(FLOPs)。一方面，它允许在固定FLOPs的情况下使用更多的通道，并增加了网络容量(从而提高了准确性)。然而，另一方面，通道数量的增加导致了更多的MAC。
按照第一节和公式1，1 × 1分组卷积的MAC和FLOPs之间的关系是：

其中g为组数，B = $hw{c}_1{c}_2/g$为FLOPs。很容易看出，给定固定的输入形状$c_1\times h\times w$和计算代价B，MAC随g的增加而增加。
为了研究实际应用中的影响，采用10个逐点分组卷积层叠加的方法构建了一个基准网络。下表报告了当总的FLOPs固定时使用不同的组数的运行速度。

很明显，使用较大的组数会显著降低运行速度。例如，在GPU上使用8组(标准密集卷积)比使用1组(标准密集卷积)慢两倍多，在ARM上慢30%。这主要是由于MAC的增加。因此，建议根据目标平台和任务，慎重选择分组号。使用较大的组数是不明智的，因为这可能会使用更多的通道，因为提高精度的好处很容易被快速增加的计算成本所抵消。

（三）、Network fragmentation reduces degree of parallelism（网络碎片化降低了并行度）

在GoogLeNet系列和auto-generated architectures中，在每个网络块中广泛采用了多路径结构。使用许多小的操作符(这里称为分段操作符)，而不是几个大的操作符。例如，在NASNET-A中，碎片运算符的数量(即一个构建块中单个卷积或池化操作的数量)是13。相反，在ResNet这样的规则结构中，这个数字是2或3。虽然这种碎片化结构有利于提高精度，但它可能会降低效率，因为它不利于GPU等具有强大并行计算能力的设备。它还引入了额外的开销，如内核启动和同步。

为了量化网络碎片化对效率的影响，作者评估了一系列具有不同碎片化程度的网络块。具体地说,构建块由1到4个1 × 1的卷积组成，这些卷积可以顺序排列，也可以并行排列。模块结构如上图。每个块重复堆叠10次。从下的结果可以看出，在GPU上fragment显著降低了速度，如4个fragment结构比1个fragment结构慢3倍。在ARM上，速度降低相对较小。

（四）、Element-wise operations are non-negligible.

如上图所示，element-wise的操作占用了相当多的时间，尤其是在GPU上。这里，element-wise的运算符包括ReLU、AddTensor、AddBias等。它们的FLOPs较小，但MAC相对较重。特别地，作者还将深度卷积作为一种逐元素运算符，因为它也具有较高的MAC/FLOPs比率。
为了验证以上结论，作者在ResNet中的“bottleneck” unit(1×1 conv后接3×3 conv后接1×1 conv，使用ReLU和shortcut连接)进行了实验。ReLU和shortcut方式操作分别被移除。下表报告了不同变体的运行时间。在去除ReLU和shortcut方式后，GPU和ARM上都获得了大约20%的加速。

基于上述准则和实证研究，得出一个有效的网络架构应该有以下几点：
1)使用平衡卷积(等信道宽度);
2)意识到使用分组卷积的代价;
3)降低破碎度;
4)减少基于元素的操作。
这些理想的属性依赖于平台特性(如内存操作和代码优化)，而这些特性超出了理论上的浮点数。在实际的网络设计中应该考虑到这些问题。
轻量化神经网络体系结构的最新进展大多基于FLOPs度量，而没有考虑以上这些特性。例如，ShuffleNet v1严重依赖于分组卷积（违背第一点）和bottleneck-like building blocks（违背第二点）。MobileNet v2使用inverted bottleneck结构（违背第而点）。它在“厚”特征图上使用深度卷积和ReLUs（违背第四点）。 auto-generated结构是高度碎片化的（违背第三点）。

三、ShuffleNet V2

（一）、ShuffleNet v1

轻量级网络的主要挑战是，在给定的计算预算(FLOPs)下，只有有限数量的特征通道是负担得起的。为了在不显著增加FLOPs的情况下增加通道数量，ShuffleNet v1中采用了两种技术:pointwise group卷积和bottleneck-like结构。然后引入通道shuffle操作，使不同通道组之间能够进行信息通信，提高精度。构建模块如下所示。

DWConv:depthwise卷积。GConv:分组卷积。
如之前节所述，点pointwise卷积和bottleneck结构都会增加MAC 。这一成本是不可忽视的，特别是轻型模型。此外，使用太多分组、在shortcut中添加元素操作也是不可取的(G4)。因此，为了实现高的模型容量和效率，关键问题是如何在不存在密集卷积和过多组的情况下，保持大量且等宽的通道。

（二）、Channel Split and ShuffleNet V2

为了达到上述目的，作者引入一个简单的算子，称为channel split。如下图所示。在每个单元开始处，将特征通道的输入c切分为 c-c’和c’两个支路。一个分支仍然作为身份映射。另一个分支由三个具有相同的输入和输出通道的卷积组成，这两个1×1卷积不再是分组卷积，部分是因为split操作而产生了两个组。

卷积后将两个分支连接起来。因此，通道的数量保持不变。然后使用channel shuffle操作来启用两个分支之间的信息通信。
shuffling之后，下一个单元开始。请注意，ShuffleNet v1中的Add操作不再存在。像ReLU和深度卷积这样的元素操作只存在于一个分支中。另外，三个连续的元素级操作Concat、Channel Shuffle和Channel Split合并为一个element-wise操作。
对于空间下采样，该单元稍作修改，如下图。删除channel split 操作符。因此，输出通道的数量增加了一倍

这些building blocks被反复堆叠起来，以构建整个网络。为了简单起见，设c ’ = c/2。总体网络结构类似于ShuffleNet v1，如下表所示。只有一个区别:在全局平均池之前添加了一个额外的1×1卷积层来混合特征，这在ShuffleNet v1中是没有的。将每个块中的通道数进行缩放，生成不同复杂性的网络，标记为0.5、1等

ShuffleNet v2不仅高效，而且准确。主要有两个原因。首先，每个构件的高效率使得使用更多的特征通道和更大的网络容量成为可能。第二，在每个块中，特征通道的一半(当c’ = c/2)直接穿过区块并加入下一个区块。这可以被看作是一种特性重用，在DenseNet和CondenseNet中有类似的操作。
在DenseNet中，为了分析特征重用模式，绘制了层间权值的l2范数，如下图所示。很明显，相邻层之间的联系比其他层更强。这意味着所有层之间的密集连接可能会引入冗余。

在ShuffleNet V2中，很容易证明第i和i+j构建块之间的“directly-connected”通道的数量是$r^{j}c$，其中r =(1−c’) / c。换句话说，特性重用的数量随着两个块之间的距离呈指数衰减。在远的块之间，特性重用变得非常弱。对于r = 0.5，如下图：

因此，ShuffleNet V2的结构通过设计实现了这类特征重用模式。与DenseNet一样，它也具有特征重用的优点，可以实现高精度，但正如前面分析的那样，它的效率要高得多。这在实验中得到了验证。

四、实验验证

消融实验是在ImageNet 2012分类数据集上进行的。所有的网络比较都有四个级别的计算复杂度，即约为40、140、300和500+ MFLOPs。这种复杂性在移动场景中很常见。其他超参数和协议与ShuffleNet v1完全相同。
1、ShuffleNet v1。g = 3时在精度和速度之间有较好的平衡。作者主要使用g = 3。

图中可看出，ShuffleNet v2模型在很大程度上优于所有其他网络，尤其是在较小的计算预算下。此外，在图像大小为224 ×224时，MobileNet v2以40 MFLOPs级别执行准确率很低，。这可能是由于通道太少。相比之下，ShuffleNet v2没有这个缺点，因为ShuffleNet v2高效设计允许使用更多的通道。此外，ShuffleNet v2和DenseNet都具有信息复用特性，但ShuffleNet v2要高效得多。
与MobileNet v1、IGCV2和IGCV3相比，虽然MobileNet v1的精度不是很好，但它在GPU上的速度比所有对手都快，包括ShuffleNet v2。这是因为它的结构满足大多数建议的指导方针(例如，对于第三点, MobileNet v1的片段甚至比ShuffleNet v2更少)。IGCV2和IGCV3速度慢。是由于使用了太多的分组卷积。这两个观察结果都符合上文提出的指导方针。
近年来，自动模型搜索（, automatic model search）已经成为CNN建筑设计中一个很有前途的趋势。它们的速度相对较慢。这主要是由于使用了太多的片段(准则三)。尽管如此，这一研究方向仍然很有前景。例如，如果将模型搜索算法与提出的准则相结合，并在目标平台上评估直接度量(速度)，就可以获得更好的模型。
与其他方法兼容。ShuffeNet v2可以与其他技术相结合，进一步提高性能。当配置了SE模块时，ShuffleNet v2的分类精度提高了0.5%，但速度有一定的损失。该块结构见下：

**推广到大模型。*虽然本文的消融实验是在轻量的情况下进行的，但ShuffleNet v2也可以用于大型模型(如FLOPs ≥2G)。下表比较了50层ShuffleNet v2与对应的ShuffleNet v1和ResNet-50。ShuffleNet v2仍然以2.3GFLOPs的表现优于ShuffleNet v1，并且以少40%的FLOPs超过ResNet-50。

模型各版本的结构如下：

对于非常深的ShuffleNet v2(例如超过100层)，为了训练更快地收敛，通过添加residual path略微修改了基本的ShuffleNet v2单元(结构如下)。下表给出了配备SE组件的164层ShuffleNet v2模型的比较。它获得了卓越的精度，超过了以前的最先进的模型与更少的FLOPs。


目标检测：为了评估对象的泛化能力，作者测试了COCO对象检测任务。使用最先进的轻型探测器Light-Head RCNN作为框架，并遵循相同的训练和测试规则。只有骨干网被取代。模型在ImageNet上进行预训练，然后在检测任务中进行微调。在训练方面，除了5000张极小集的图像外，在COCO中使用了训练+val集，并使用极小集进行测试。精度度量是COCO标准mmAP，即从0.5到0.95的box IoU阈值处的平均map。实验结果如下：

将ShuffleNet v2与Xception、ShuffleNet v1和MobileNet v2这三种轻量级模型在四个层次上进行比较。上表结果显示，ShuffleNet v2的性能最好。
比较之前的检测结果和分类结果，分类上的准确率排名为ShuffleNet v2 ≥MobileNet v2 >ShuffeNet v1> Xception，检测时却变为ShuffleNet v2 >Xception≥ShuffleNet v1≥ MobileNet v2。这表明Xception在检测任务上很好。这可能是由于Xception构建块比其他同类更大的感受野(7 vs. 3)。受此启发，作者还通过在每个构建块的第一个pointwise卷积之前引入额外的3× 3pointwise卷积来扩大ShuffleNet v2的感受野。此变体表示为ShuffleNet v2。它进一步提高了准确性。
作者还在GPU上测试了运行时间。为了公平比较，batch size设置为4，以确保充分的GPU利用率。由于数据复制(分辨率高达800 × 1200)和其他操作(如PSRoI Pooling)的开销，不同模型之间的速度差距小于分类的速度差距。尽管如此，ShuffleNet v2仍然优于其他的，例如，比ShuffleNet v1快40%左右，比MobileNet v2快16%。
此外，改型ShuffleNet v2*具有最好的精度，并且仍然比其他方法更快。这激发了一个实际的问题:如何增加感受野的大小?这对于高分辨率图像中的目标检测至关重要。
源码如下：

import torch
from torch import Tensor
import torch.nn as nn
from typing import Callable, Any, List
import torch.utils.model_zoo as model_zoo

__all__ = [
    'ShuffleNetV2', 'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0',
    'shufflenet_v2_x1_5', 'shufflenet_v2_x2_0'
]

model_urls = {
    'shufflenetv2_x0.5': 'https://download.pytorch.org/models/shufflenetv2_x0.5-f707e7126e.pth',
    'shufflenetv2_x1.0': 'https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth',
    'shufflenetv2_x1.5': None,
    'shufflenetv2_x2.0': None,
}


def channel_shuffle(x: Tensor, groups: int) -> Tensor:
    batchsize, num_channels, height, width = x.size()
    channels_per_group = num_channels // groups

    # reshape
    x = x.view(batchsize, groups,
               channels_per_group, height, width)

    x = torch.transpose(x, 1, 2).contiguous()

    # flatten
    x = x.view(batchsize, -1, height, width)

    return x


def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class BasicBlock(nn.Module):
    expansion = 1
    __constants__ = ['downsample']

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None):
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    expansion = 4
    __constants__ = ['downsample']

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups

        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class InvertedResidual(nn.Module):
    def __init__(
            self,
            inp: int,
            oup: int,
            stride: int
    ) -> None:
        super(InvertedResidual, self).__init__()

        if not (1 <= stride <= 3):
            raise ValueError('illegal stride value')
        self.stride = stride

        branch_features = oup // 2
        assert (self.stride != 1) or (inp == branch_features << 1)

        if self.stride > 1:
            self.branch1 = nn.Sequential(
                self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1),
                nn.BatchNorm2d(inp),
                nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(branch_features),
                nn.ReLU(inplace=True),
            )
        else:
            self.branch1 = nn.Sequential()

        self.branch2 = nn.Sequential(
            nn.Conv2d(inp if (self.stride > 1) else branch_features,
                      branch_features, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(branch_features),
            nn.ReLU(inplace=True),
            self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1),
            nn.BatchNorm2d(branch_features),
            nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(branch_features),
            nn.ReLU(inplace=True),
        )

    @staticmethod
    def depthwise_conv(
            i: int,
            o: int,
            kernel_size: int,
            stride: int = 1,
            padding: int = 0,
            bias: bool = False
    ) -> nn.Conv2d:
        return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)

    def forward(self, x: Tensor) -> Tensor:
        if self.stride == 1:
            x1, x2 = x.chunk(2, dim=1)
            out = torch.cat((x1, self.branch2(x2)), dim=1)
        else:
            out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)

        out = channel_shuffle(out, 2)

        return out


class ShuffleNetV2(nn.Module):
    def __init__(
            self,
            stages_repeats: List[int],
            stages_out_channels: List[int],
            num_classes: int = 1000,
            inverted_residual: Callable[..., nn.Module] = InvertedResidual
    ) -> None:
        super(ShuffleNetV2, self).__init__()

        if len(stages_repeats) != 3:
            raise ValueError('expected stages_repeats as list of 3 positive ints')
        if len(stages_out_channels) != 5:
            raise ValueError('expected stages_out_channels as list of 5 positive ints')
        self._stage_out_channels = stages_out_channels

        input_channels = 3
        output_channels = self._stage_out_channels[0]
        self.conv1 = nn.Sequential(
            nn.Conv2d(input_channels, output_channels, 3, 2, 1, bias=False),
            nn.BatchNorm2d(output_channels),
            nn.ReLU(inplace=True),
        )
        input_channels = output_channels

        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        # Static annotations for mypy
        self.stage2: nn.Sequential
        self.stage3: nn.Sequential
        self.stage4: nn.Sequential
        stage_names = ['stage{}'.format(i) for i in [2, 3, 4]]
        for name, repeats, output_channels in zip(
                stage_names, stages_repeats, self._stage_out_channels[1:]):
            seq = [inverted_residual(input_channels, output_channels, 2)]
            for i in range(repeats - 1):
                seq.append(inverted_residual(output_channels, output_channels, 1))
            setattr(self, name, nn.Sequential(*seq))
            input_channels = output_channels

        output_channels = self._stage_out_channels[-1]
        self.conv5 = nn.Sequential(
            nn.Conv2d(input_channels, output_channels, 1, 1, 0, bias=False),
            nn.BatchNorm2d(output_channels),
            nn.ReLU(inplace=True),
        )
        self.out_features = output_channels
        self.bottleneck = nn.Sequential(
            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.5)
        )
        self.bottleneck[0].weight.data.normal_(0, 0.005)
        self.bottleneck[0].bias.data.fill_(0.1)
        self.head = nn.Sequential(
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )
        # self.fc = nn.Linear(output_channels, num_classes)

    def _forward_impl(self, x: Tensor) -> Tensor:
        # See note [TorchScript super()]
        x = self.conv1(x)
        x = self.maxpool(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = self.conv5(x)
        x = x.mean([2, 3])  # globalpool
        # print("x.shape={}".format(x.shape))x.shape=torch.Size([48, 1024])
        # x = self.fc(x)
        # print("x2.size={}".format(x.shape))
        x = self.bottleneck(x)
        x = self.head(x)
        return x

    def forward(self, x: Tensor) -> Tensor:
        return self._forward_impl(x)


def _shufflenetv2(arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs: Any) -> ShuffleNetV2:
    model = ShuffleNetV2(*args, **kwargs)

    if pretrained:
        model_url = model_urls[arch]
        if model_url is None:
            raise NotImplementedError('pretrained {} is not supported as of now'.format(arch))
        else:
            state_dict = model_zoo.load_url(model_url, model_dir='.')
            model.load_state_dict(state_dict, strict=False)
            # state_dict = load_state_dict_from_url(model_url, progress=progress)
            # model.load_state_dict(state_dict, False)

    return model


def shufflenet_v2_x0_5(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2:
    """
    Constructs a ShuffleNetV2 with 0.5x output channels, as described in
    `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
    <https://arxiv.org/abs/1807.11164>`_.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _shufflenetv2('shufflenetv2_x0.5', pretrained, progress,
                         [4, 8, 4], [24, 48, 96, 192, 1024], **kwargs)


def shufflenet_v2_x1_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2:
    """
    Constructs a ShuffleNetV2 with 1.0x output channels, as described in
    `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
    <https://arxiv.org/abs/1807.11164>`_.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _shufflenetv2('shufflenetv2_x1.0', pretrained, progress,
                         [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs)


def shufflenet_v2_x1_5(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2:
    """
    Constructs a ShuffleNetV2 with 1.5x output channels, as described in
    `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
    <https://arxiv.org/abs/1807.11164>`_.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _shufflenetv2('shufflenetv2_x1.5', pretrained, progress,
                         [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs)


def shufflenet_v2_x2_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2:
    """
    Constructs a ShuffleNetV2 with 2.0x output channels, as described in
    `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design"
    <https://arxiv.org/abs/1807.11164>`_.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _shufflenetv2('shufflenetv2_x2.0', pretrained, progress,
                         [4, 8, 4], [24, 244, 488, 976, 2048], **kwargs)