Classification基础实验系列四——MobileNet v2论文笔记与复现

一、概述

上一篇博客记录了MobileNet v1的学习过程，了解到该网络主要是由深度可分离卷积（depthwise + pointwise）“直筒状”拼接而成的简单神经网络。本博客为v2的学习笔记。为了切入重点，我们先来分析下MobileNet v1主要存在哪些缺点：

未使用Short cut
ResNet, DenseNet等结构早已证明，Short cut(或者skip connection)可以通过特征复用将多尺度的特征融合在一起，进而极大提升网络的性价比。而MobileNet v1由于年代较久远，还停留在“直筒状”结构。
Depthwise与Non-linearity不搭
ShuffleNet和Xception中都建议depthwise conv后面不要使用non-linearity，在Xception中还有实验对比：

Fig. 1. depthwise不适合于non-linearity一同食用

论文中提到的一个合理解释是：depth很关键，由于Inception中的中间层是全channel的，所以non-linearity是有利的，但是对于depthwise conv，其特征是单channle的，所以non-linearity有害，可能造成了信息丢失。

MobileNet v2就针对这两个问题作出了相应的改进。文章花了大量篇幅论述ReLU(也代表其他non-linearity)在对通道数较少（维度较低）的卷积层输出进行操作时，会造成信息的损失。因此如何避免ReLU的这一特性成为本文的核心改进方向。

相比MobileNet v1， v2主要引入了两个改进：Linear Bottleneck和Inverted Residual Block。
v1和v2主要模块的对比：（图片出自该知乎文章）。

Fig. 2. MobileNet v1与v2对比

二、论文笔记

2.1 Manifold of interest

"Manifold of interest"，可以翻译成兴趣流形，是本文中最重要的概念之一。我的理解是，对于m个输入样本而言，网络每一层L_i的输出张量（注意是应用激活函数之前的张量）可以表述成的体素，其中d_i为该体素的深度，或者说维度（dimensionality）。我们只拿其中一层L_i来说，m个输入图像的信息在该层被映射为一个的tensor。而该tensor中并不是每一个像素对于表征输入图像而言都是不可或缺的。可能只有一部分像素，就足够表征这些输入图像在该层的某种感兴趣信息（比如前几层可能是m个图像对应的轮廓，后几层是对应的眼睛鼻子等组件）。这些真正“不可或缺”的像素，在的特征空间中呈一个流形(manifold)分布。，我们说这个兴趣流形是嵌入在特征空间的低维子空间(low-dimensional subspace)中的。

这个理论听起来似乎很有用，如果兴趣流形只存在于低维子空间中，那就可以对卷积层的输出特征进行降维，理想的情况是我们可以一直降维到该流形“张满”的那个特定维度对应的子空间。这样的显然可以节省大量的参数。该思想在MobileNet v1中得到了应用（weight multiplier）并取得了不错效果。

然而这个理论有一个严重的问题，就是以ReLU为代表的non-linearity会严重破坏输入空间的信息。举一个最简单的例子，1维空间中的直线经过ReLU会变成射线。而n维空间中的曲面经过ReLU可能只剩下具有n个节点的分段线性曲线。另一方面，从数值的角度说，ReLU会使激活feature map变稀疏。极端的情况，假如某一维（一个通道）的输出均为负数，那么通过ReLU之后，输出tensor相当于被降维了。

作者在Figure 1部分详细论述了这个问题：将一个二维的螺旋线（代表2D空间中的1D流形）通过随机矩阵T（后接RELU）转换至n-维tensor；再通过将其投影回2D空间。如果在转换过程中没有信息损失（比如去掉ReLU成为线性变换），则得到的2D矩阵应该和原来一样。然而由于ReLU的作用，当n较小时，恢复出来的图像出现了明显的“崩塌”，即流形中的很多点和其他点重叠到了一起。然而当n变大到15或30时，信息丢失会逐渐减少。这里的“崩塌(collapse)”可以理解成上一段中由于通道中所有数值为负而被ReLU降维这一现象。

Fig. 3. 若特征空间维度太低，ReLU会造成流形的信息丢失

根据上述分析以及实验，作者得到如下结论：如果每层卷积输出的tensor具有足够多的通道数，在经过ReLU之后，即便一些通道崩塌了，其他通道可能仍然保留了足够的信息。作者在论文后面给出了证明：如果输入流形可以被嵌入到activation space的维度足够低的子空间中（换句话说就是特征空间维度足够高），那么ReLU就可以在保留信息的同时完成non-linearity的本职工作——提高模型的表达能力。所以问题就变得清晰了——在必要的ReLU之前，提高卷积层输出tensor的维度。

2.2 Linear Bottleneck

经过上面的分析，现在可以很自然地引出改论文的第一个改进——Linear Bottleneck。Bottleneck这个词我也没找到严格的定义，感觉通常可以简单理解成维度较低的输出特征前面的层（比如这里depthwise之后的1x1卷积输出的特征即为bottleneck feature，所以每个block最后一个1x1卷积层即为Bottleneck）。所以Linear Bottleneck就是在1x1 conv之后去掉了non-linearity。之所以去掉这一层的ReLU也是因为Bottleneck比较薄，如果用ReLU的话就会出现上面的问题。

2.3 Expansion Convolution Layer

这一改进也是旨在解决同样的问题，即作者为了避免输入ReLU之前的特征太“薄”，特意先给输入特征升维。这里的升维也是通过1x1 Conv。这样一来，整个block从外观上看就是两头薄，中间厚，故该1x1 Conv被称为Expansion Conv。

此外，作者在该Block基础上又加上了short cut。其作用在于：①即便输入特征深度为0，也能利用short cut将卷积层转换为恒等映射。②利用short cut的通用功能，即特征融合。由于传统的残差块结构是中间薄，两头厚，所以作者将v2的这个1x1 Expansion Conv - ReLU6 - Seperable Conv - ReLU6 - Linear Bottleneck结构称为Inverted Residual Block。图片出自该知乎文章。

Fig. 4. 对比传统的残差块

一个小问题是，能不能不要ReLU，也不要Expansion Conv，都用Linear Conv？答案显然是不行的，作为提取特征的层必须有足够的表达能力，一个由纯linear layer组成的网络再深也只相当于一层。

2.4 网络架构

值得注意的是，虽然MobileNet v2的单个block比v1参数要多，然而整体参数是比v1要少的。对比下网络架构就可以看出——相比一般网络架构的堆叠方式，v2看起来没那么“规整”，比如32维特征后面是24维，而不是传统的32或者64。

Fig. 5. 最终网络架构

三、实现

3.1 网络搭建

上一篇中我详细了解了MobileNet中depthwise的Gluon实现，清晰起见写的比较啰嗦。其实depthwise不需要专门定义一个nn.Block，直接用Conv2D即可。这里我重新整理了一下相关的实现，得到一个相对简洁的模型脚本：

import mxnet as mx
from mxnet.gluon import nn
from mxnet import nd, autograd
from mxnet.gluon import data as gdata
from mxnet.gluon.model_zoo import vision

class ReLU6(nn.HybridBlock):
    def __init__(self, **kwags):
        super(ReLU6, self).__init__(**kwags)
    
    def hybrid_forward(self, F, x):
        return F.clip(x, 0, 6)

def ConvBlock(channels, kernel_size, strides, padding=1, groups=1, activation='relu6'):
    block = nn.HybridSequential()
    block.add(nn.Conv2D(channels, kernel_size, strides, padding=padding, groups=groups, use_bias=False))
    block.add(nn.BatchNorm())
    if activation is not None:
        block.add(ReLU6(prefix='relu6_'))
    return block

def DepthWiseConv(channels, strides):
    return ConvBlock(channels, 3, strides, groups=channels)
    
def LinearBottleneck(channels):
    return ConvBlock(channels, 1, 1, 0, activation=None)

def ExpansionConv(channels):
    return ConvBlock(channels, 1, 1, 0)

class InvertedResidual(nn.HybridBlock):
    def __init__(self, in_channels, out_channels, strides, t=6, **kwags):
        super(InvertedResidual,self).__init__(**kwags)
        self.strides = strides
        self.keep_channels = in_channels == out_channels
        expanded_channels = t * in_channels
        self.inver_residual = nn.HybridSequential()
        with self.inver_residual.name_scope():
            self.inver_residual.add(ExpansionConv(expanded_channels),
                                DepthWiseConv(expanded_channels, strides),
                                LinearBottleneck(out_channels))
    
    def hybrid_forward(self, F, x):
        out = self.inver_residual(x)
        if self.strides == 1 and self.keep_channels:
            out = out + x
            #out = F.elemwise_add(out, x)
        return out
    
def RepeatedInvertedResiduals(in_channels, out_channels, repeats, strides, t, **kwags):
    sequence = nn.HybridSequential(**kwags)
    # The first layer of each sequence has a stride s and all others use stride 1.
    sequence.add(InvertedResidual(in_channels, out_channels, strides, t))
    for _ in range(1, repeats):
        sequence.add(InvertedResidual(out_channels, out_channels, 1, t))
    return sequence

class MobileNetV2(nn.HybridBlock):
    def __init__(self, num_classes, width_multiplier=1.0, **kwags):
        super(MobileNetV2, self).__init__(**kwags)
        input_feature_channels = int(32 * width_multiplier)
        
        self.bottleneck_settings = [
            # t, c, n, s
            [1, 16, 1, 1, "stage0_"],      # -> 112x112
            [6, 24, 2, 2, "stage1_"],      # -> 56x56
            [6, 32, 3, 2, "stage2_"],      # -> 28x28
            [6, 64, 4, 2, "stage3_0_"],    # -> 14x14
            [6, 96, 3, 1, "stage3_1_"],    # -> 14x14
            [6, 160, 3, 2, "stage4_0_"],   # -> 7x7
            [6, 320, 1, 1, "stage4_1_"],   # -> 7x7
        ]
        self.net = nn.HybridSequential()
        self.net.add(ConvBlock(input_feature_channels, 3, 2))
        
        in_channels = input_feature_channels
        for t, c, n, s, prefix in self.bottleneck_settings:
            out_channels = int(width_multiplier * c)
            self.net.add(RepeatedInvertedResiduals(in_channels, out_channels, n, s, t, prefix=prefix))
            in_channels = out_channels  # 下一层的输入通道数为当前层的输出通道数
        
        # 注意：MobileNetV2使用的分类头不是GAP + Dense，而是GAP + 1x1 Linear Conv + Flatten
        self.net.add(ConvBlock(int(1280*width_multiplier), 1, 1, 0),
                          nn.GlobalAvgPool2D(),
                          nn.Conv2D(num_classes, 1, 1, 0, activation=None, use_bias=False),
                          nn.Flatten())
    
    def hybrid_forward(self, F, x):
        return self.net(x)

值得注意的在堆叠Inverted Residual Block时，里面有一些小细节论文中可能没有详细说明，或者很容易被忽略掉：
- Fig. 5 中重复6次的bottleneck，只有第一次strides为s，其余strides为1。（否则会改变空间维度，也没法重复...）
- 关于什么时候用short cut：注意，论文中shortcut的连接方式是element wise add。既然是逐元素相加，这就要求：相加的两个特征①空间维度相同；②通道数一样多。如果是concat就不需要②。根据这两个要求，strides不为1的Bottleneck，shortcut不能使用；由于V2中的shortcut旁支不负责改变维度（通道数），所以对于输入通道数不等于输出通道数的Bottleneck，shortcut不能使用；
- 我看到有些人的代码中，最后一个1x1 conv，即1280那个不乘以width_multiplier。但是论文中我没找到对应的地方，就先按所有通道都乘缩放因子来写。
  
  Fig. 6. MobileNet v2的shortcut不是每个block都使用