pytorch实现并训练MobileNetV3

1.MobileNetV3

MobileNet是Google公司推出的轻量化系列网络，用以在移动平台上进行神经网络的部署和应用。2019年，Google发布了第三代MobileNet，即MobileNetV3。在MobileNet系列的精度和计算量上都达到了新的state-of-art，以下简单回顾一下三代MobileNet的主要特性：

1. MobileNetV1:MobileNetV1的主要思想是将普通的卷积操作分解为两步，先做一次仅卷积，不求和的“深度卷积”操作，再使用1*1的“点卷积”对深度卷积得到的多通道结果进行融合，减少了大量的通道融合时间和参数。
2. MobileNetV2:ResNet中的Bottleneck残差块走下来是通道数先变少在变多，每层中的激活函数会导致丢失一部分的信息（神经元抑制），那么在输入通道数很多时丢的信息会很多，因此MobileNetV2把Bottleneck改为先将通道数增加，再将通道数减少，引入到MobileNetV2中，称之为“反转残差块”。
3. MobileNetV3：MobileNetV3主要的贡献就是使用科学的理论重新设计了网络结构，并且引入了H-Swish激活函数与Relu搭配使用。另外在网络中还引入了Squeeze-And-Excite模块，这个虽然在文中只是简单的一提，但是却很重要。复现的时候没有这个是达不到论文的效果的。

2.网络结构

论文中根据硬件条件的不同设计了两套网络，分别是MobileNetV3_large和MobileNetV3_small，前者适合硬件条件较好的设备，后者适合硬件条件局限性较大的设备，精度略有下降。网络结构如下图：

1. MobileNetV3_large:
   

2. MobileNetV3_small:
   
   下面的注释说明都很重要，比如Squeeze-And-Excite模块，在复现的时候都不可以漏掉。

3.pytorch实现

网络的实现比较简单，但网络上的几个pytorch复现版本有的没有严格按照论文描述的执行，有的有错误，处于学习的目的，我按照论文里描述的网络结构，借鉴了部分网络代码，使用pytorch框架复现了MobileNetV3_large和MobileNetV3_small，自己使用oxFlower17分类的数据集训练了一下，确保是可行的。完整的项目链接为：https://github.com/yichaojie/MobileNetV3

项目中model下的model.py为主要网络结构。主要的网络结构代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F

def Hswish(x,inplace=True):
    return x * F.relu6(x + 3., inplace=inplace) / 6.

def Hsigmoid(x,inplace=True):
    return F.relu6(x + 3., inplace=inplace) / 6.


# Squeeze-And-Excite模块
class SEModule(nn.Module):
    def __init__(self, channel, reduction=4):
        super(SEModule, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.se = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y=self.avg_pool(x).view(b, c)
        y=self.se(y)
        y = Hsigmoid(y).view(b, c, 1, 1)
        return x * y.expand_as(x)

class Bottleneck(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,exp_channels,stride,se='True',nl='HS'):
        super(Bottleneck, self).__init__()
        padding = (kernel_size - 1) // 2
        if nl == 'RE':
            self.nlin_layer = F.relu6
        elif nl == 'HS':
            self.nlin_layer = Hswish
        self.stride=stride
        if se:
            self.se=SEModule(exp_channels)
        else:
            self.se=None
        self.conv1=nn.Conv2d(in_channels,exp_channels,kernel_size=1,stride=1,padding=0,bias=False)
        self.bn1 = nn.BatchNorm2d(exp_channels)
        self.conv2=nn.Conv2d(exp_channels,exp_channels,kernel_size=kernel_size,stride=stride,
                             padding=padding,groups=exp_channels,bias=False)
        self.bn2=nn.BatchNorm2d(exp_channels)
        self.conv3=nn.Conv2d(exp_channels,out_channels,kernel_size=1,stride=1,padding=0,bias=False)
        self.bn3=nn.BatchNorm2d(out_channels)
        # 先初始化一个空序列，之后改造其成为残差链接
        self.shortcut = nn.Sequential()
        # 只有步长为1且输入输出通道不相同时才采用跳跃连接(想一下跳跃链接的过程，输入输出通道相同这个跳跃连接就没意义了)
        if stride == 1 and in_channels != out_channels:
            self.shortcut = nn.Sequential(
                # 下面的操作卷积不改变尺寸，仅匹配通道数
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self,x):
        out=self.nlin_layer(self.bn1(self.conv1(x)))
        if self.se is not None:
            out=self.bn2(self.conv2(out))
            out=self.nlin_layer(self.se(out))
        else:
            out = self.nlin_layer(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out = out + self.shortcut(x) if self.stride == 1 else out
        return out


class MobileNetV3_large(nn.Module):
    # (out_channels,kernel_size,exp_channels,stride,se,nl)
    cfg=[
        (16,3,16,1,False,'RE'),
        (24,3,64,2,False,'RE'),
        (24,3,72,1,False,'RE'),
        (40,5,72,2,True,'RE'),
        (40,5,120,1,True,'RE'),
        (40,5,120,1,True,'RE'),
        (80,3,240,2,False,'HS'),
        (80,3,200,1,False,'HS'),
        (80,3,184,1,False,'HS'),
        (80,3,184,1,False,'HS'),
        (112,3,480,1,True,'HS'),
        (112,3,672,1,True,'HS'),
        (160,5,672,2,True,'HS'),
        (160,5,960,1,True,'HS'),
        (160,5,960,1,True,'HS')
    ]
    def __init__(self,num_classes=17):
        super(MobileNetV3_large,self).__init__()
        self.conv1=nn.Conv2d(3,16,3,2,padding=1,bias=False)
        self.bn1=nn.BatchNorm2d(16)
        # 根据cfg数组自动生成所有的Bottleneck层
        self.layers = self._make_layers(in_channels=16)
        self.conv2=nn.Conv2d(160,960,1,stride=1,bias=False)
        self.bn2=nn.BatchNorm2d(960)
        # 卷积后不跟BN，就应该把bias设置为True
        self.conv3=nn.Conv2d(960,1280,1,1,padding=0,bias=True)
        self.conv4=nn.Conv2d(1280,num_classes,1,stride=1,padding=0,bias=True)

    def _make_layers(self,in_channels):
        layers=[]
        for out_channels,kernel_size,exp_channels,stride,se,nl in self.cfg:
            layers.append(
                Bottleneck(in_channels,out_channels,kernel_size,exp_channels,stride,se,nl)
            )
            in_channels=out_channels
        return nn.Sequential(*layers)

    def forward(self,x):
        out=Hswish(self.bn1(self.conv1(x)))
        out=self.layers(out)
        out=Hswish(self.bn2(self.conv2(out)))
        out=F.avg_pool2d(out,7)
        out=Hswish(self.conv3(out))
        out=self.conv4(out)
        # 因为原论文中最后一层是卷积层来实现全连接的效果，维度是四维的，后两维是1，在计算损失函数的时候要求二维，因此在这里需要做一个resize
        a,b=out.size(0),out.size(1)
        out=out.view(a,b)
        return out

class MobileNetV3_small(nn.Module):
    # (out_channels,kernel_size,exp_channels,stride,se,nl)
    cfg = [
        (16,3,16,2,True,'RE'),
        (24,3,72,2,False,'RE'),
        (24,3,88,1,False,'RE'),
        (40,5,96,2,True,'HS'),
        (40,5,240,1,True,'HS'),
        (40,5,240,1,True,'HS'),
        (48,5,120,1,True,'HS'),
        (48,5,144,1,True,'HS'),
        (96,5,288,2,True,'HS'),
        (96,5,576,1,True,'HS'),
        (96,5,576,1,True,'HS')
    ]
    def __init__(self,num_classes=17):
        super(MobileNetV3_small,self).__init__()
        self.conv1=nn.Conv2d(3,16,3,2,padding=1,bias=False)
        self.bn1=nn.BatchNorm2d(16)
        # 根据cfg数组自动生成所有的Bottleneck层
        self.layers = self._make_layers(in_channels=16)
        self.conv2=nn.Conv2d(96,576,1,stride=1,bias=False)
        self.bn2=nn.BatchNorm2d(576)
        # 卷积后不跟BN，就应该把bias设置为True
        self.conv3=nn.Conv2d(576,1280,1,1,padding=0,bias=True)
        self.conv4=nn.Conv2d(1280,num_classes,1,stride=1,padding=0,bias=True)

    def _make_layers(self,in_channels):
        layers=[]
        for out_channels,kernel_size,exp_channels,stride,se,nl in self.cfg:
            layers.append(
                Bottleneck(in_channels,out_channels,kernel_size,exp_channels,stride,se,nl)
            )
            in_channels=out_channels
        return nn.Sequential(*layers)

    def forward(self,x):
        out=Hswish(self.bn1(self.conv1(x)))
        out=self.layers(out)
        out=self.bn2(self.conv2(out))
        se=SEModule(out.size(1))
        out=Hswish(se(out))
        out = F.avg_pool2d(out, 7)
        out = Hswish(self.conv3(out))
        out = self.conv4(out)
        # 因为原论文中最后一层是卷积层来实现全连接的效果，维度是四维的，后两维是1，在计算损失函数的时候要求二维，因此在这里需要做一个resize
        a, b = out.size(0), out.size(1)
        out = out.view(a, b)
        return out

# 测试代码，跑通证明网络结构没问题
# def test():
#     net=MobileNetV3_small()
#     x=torch.randn(2,3,224,224)
#     y=net(x)
#     print(y.size())
#     print(y)
#
# if __name__=="__main__":
#     test()

我是用牛津花朵17分类数据集进行了分类，运行train.py即可开始训练，batch-size大小默认为64，在MobileNetV3_large上测试训练32个epoch左右即可收敛：

在inference.py中可以选择权重和图片来进行推断测试。

4.训练自定义数据集

4.1划分数据集

首先划分好自己的数据集，例如有1000类，则生成名为0-999的文件夹，把对应同一类别的照片放到同一文件夹下，可以划分出训练集，测试机和验证集，格式都是这样。然后将训练集放到data/splitData/train下，测试集放到data/splitData/test下，验证集放到data/splitData/valid下即可。具体的可以参照我项目中的格式。

4.2开始训练

在train.py的头部可以更改batch-size，epoch、lr等参数：

在train.py的55行可以选择要训练的模型，切换MobileNetV3_large和MobileNetV3_small：

更改完成后运行train.py即可开始训练，训练的权重会保存在weights目录下，best.pkl代表最好结果，last.pkl代表最后一个epoch得到的权重。