MicroNet关键代码解读（Micro-block与Dynamic Shift-Max的实现代码）

论文地址：https://arxiv.org/pdf/2011.12289
中文翻译：https://hpg123.blog.csdn.net/article/details/141772832?spm=1001.2014.3001.5502
发表时间：2022
项目地址：https://github.com/liyunsheng13/micronet

在MicroNet论文中提出了Micro-block与Dynamic Shift-Max，这里对其代码实现进行深入分析。最终结论是Micro-block的定义实现十分混乱，MicroNet比Moblienet强主要是深度分离卷积的性能没有充分挖掘到位，可以替换成conv1xk_group+convkx1_group的组合，从而在低flop的约束条件下实现了5个点左右的提升；另外一点是使用了带参数的激活函数，同时激活函数中提供了group间的数据交互，故再次提升了模型精度。

1、Micro-block

1.1 模块定义

在MicroNet论文中一共有Micro-Block-A、Micro-Block-B、Micro-Block-C三种模块设计。

Micro-Block-A 采用微因式点卷积和深度卷积的精简组合（见图 2-右），它主要适用在模型的浅层。需要注意的是，微因数深度卷积扩大了通道数，而分组自适应卷积则压缩了通道数。

Micro-Block-B 用于连接 Micro-Block-A 和 Micro-Block-C。与 Micro-Block-A 不同的是，它使用的是全 Micro-Factorized 点式卷积，其中包括两个组自适应卷积（如图 5b 所示）。前者压缩了通道数，后者则扩大了通道数。

微块-C（如图 5c 所示）使用了常规组合。 深度卷积和点卷积的正则组合。它用于模型的深层，因为与精简组合相比，它在通道融合（点上）上花费的计算量更大。

基于下列图片可以发现，Micro-block虽然使用了group conv，但基于Φ函数的shitf操作，是可以实现数据在不同group间的交互。

1.2 模块效果

首先通过论文中的表1可以发现Micro-Block-B 是一直只有一个，且介于A与C之间。

在block级别的对比中，可以发现Micro-block具备明显的涨点效果。

在与MobileNet的对比中可以发现在没有Shift-Max操作时，Micro就可以比Mobile高6个点了。这主要是因为MobileNet只是将conv拆解为深度分离卷积与点卷积，conv3x3的操作没有拆解；而在Micro中，将conv拆解为了conv3x1_group+full_group_connect+conv1x3_group，这种拆解相比MobileNet更高效。同时在保证相同flop的时候能具备更深的网络结构。

1.3 代码实现

以下代码推测是Micro-Block-A 的实现，可以发现是 kx1+bn+1xk+ChannelShuffle操作，对比DepthSpatialSepConv函数，发现实现上高度接近。

class ChannelShuffle(nn.Module):
    def __init__(self, groups):
        super(ChannelShuffle2, self).__init__()
        self.groups = groups

    def forward(self, x):
        b, c, h, w = x.size()

        channels_per_group = c // self.groups

        # reshape
        x = x.view(b, self.groups, channels_per_group, h, w)

        x = torch.transpose(x, 1, 2).contiguous()
        out = x.view(b, -1, h, w)

        return out

######################################################################3
# part 3: new block
#####################################################################3

class SpatialSepConvSF(nn.Module):
    def __init__(self, inp, oups, kernel_size, stride):
        super(SpatialSepConvSF, self).__init__()

        oup1, oup2 = oups
        self.conv = nn.Sequential(
            nn.Conv2d(inp, oup1,
                (kernel_size, 1),
                (stride, 1),
                (kernel_size//2, 0),
                bias=False, groups=1
            ),
            nn.BatchNorm2d(oup1),
            nn.Conv2d(oup1, oup1*oup2,
                (1, kernel_size),
                (1, stride),
                (0, kernel_size//2),
                bias=False, groups=oup1
            ),
            nn.BatchNorm2d(oup1*oup2),
            ChannelShuffle(oup1),
        )

    def forward(self, x):
        out = self.conv(x)
        return out

在追溯DYMicroBlock代码中并为明确发现Micro-Block-A、Micro-Block-B、Micro-Block-C的定义。但发现了不少ChannelShuffle操作，这表明在conv_group模型中，通道间的交互是不可少的。

2、Dynamic Shift-Max

2.1 模块定义

在MicroNet中提到了一种新的激活函数–动态 Shift-Max函数，以增强非线性。它能动态地将输入特征图与它的环形组移动进行动态融合，以group为移动单位进行移动。Dynamic Shift-Max 还能
加强组之间的联系。这是对Micro-Factorized pointwise convolution的补充，侧重于组内连接的互补性。

其具体作用如图4所示，对channel以group为单位进行循环移动，并基于fc映射后动态输出max值。

2.2 模块效果

Shift-Max模块的效果如下所示，可以看到相比于不加之前，top1 acc提升了2.7%，而当使用dynamic shift-max后，top1 acc相比于不加提升了6.8%。可以发现Shift-Max模块是MicroNet的涨点关键。其对于分组卷积提供了group间的特征互动，同时相比于普通的激活函数，其是带可训练参数的模型。

同时在与其他同类型激活函数对比中，可以发现使用Dynamic Shift-Max时更加有效；同时Dynamic ReLU也证明了在低flop模型中也具有显著的作用，只是在group conv中略有不殆。

2.3 代码实现

class DYShiftMax(nn.Module):
    def __init__(self, inp, oup, reduction=4, act_max=1.0, act_relu=True, init_a=[0.0, 0.0], init_b=[0.0, 0.0], relu_before_pool=False, g=None, expansion=False):
        super(DYShiftMax, self).__init__()
        self.oup = oup
        self.act_max = act_max * 2
        self.act_relu = act_relu
        self.avg_pool = nn.Sequential(
                nn.ReLU(inplace=True) if relu_before_pool == True else nn.Sequential(),
                nn.AdaptiveAvgPool2d(1)
            )

        self.exp = 4 if act_relu else 2
        self.init_a = init_a
        self.init_b = init_b

        # determine squeeze
        squeeze = _make_divisible(inp // reduction, 4)
        if squeeze < 4:
            squeeze = 4
        print('reduction: {}, squeeze: {}/{}'.format(reduction, inp, squeeze))
        print('init-a: {}, init-b: {}'.format(init_a, init_b))

        self.fc = nn.Sequential(
                nn.Linear(inp, squeeze),
                nn.ReLU(inplace=True),
                nn.Linear(squeeze, oup*self.exp),
                h_sigmoid()
        )
        if g is None:
            g = 1
        self.g = g[1]
        if self.g !=1  and expansion:
            self.g = inp // self.g
        print('group shuffle: {}, divide group: {}'.format(self.g, expansion))
        self.gc = inp//self.g
        index=torch.Tensor(range(inp)).view(1,inp,1,1)
        index=index.view(1,self.g,self.gc,1,1)
        indexgs = torch.split(index, [1, self.g-1], dim=1)
        indexgs = torch.cat((indexgs[1], indexgs[0]), dim=1)
        indexs = torch.split(indexgs, [1, self.gc-1], dim=2)
        indexs = torch.cat((indexs[1], indexs[0]), dim=2)
        self.index = indexs.view(inp).type(torch.LongTensor)
        self.expansion = expansion

    def forward(self, x):
        x_in = x
        x_out = x

        b, c, _, _ = x_in.size()
        y = self.avg_pool(x_in).view(b, c)
        y = self.fc(y).view(b, self.oup*self.exp, 1, 1)
        y = (y-0.5) * self.act_max

        n2, c2, h2, w2 = x_out.size()
        x2 = x_out[:,self.index,:,:]

        if self.exp == 4:
            a1, b1, a2, b2 = torch.split(y, self.oup, dim=1)

            a1 = a1 + self.init_a[0]
            a2 = a2 + self.init_a[1]

            b1 = b1 + self.init_b[0]
            b2 = b2 + self.init_b[1]

            z1 = x_out * a1 + x2 * b1
            z2 = x_out * a2 + x2 * b2

            out = torch.max(z1, z2)

        elif self.exp == 2:
            a1, b1 = torch.split(y, self.oup, dim=1)
            a1 = a1 + self.init_a[0]
            b1 = b1 + self.init_b[0]
            out = x_out * a1 + x2 * b1

        return out