Pytorch MobileNetV1 学习

"""mobilenet in pytorch



[1] Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, Hartwig Adam

    MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
    https://arxiv.org/abs/1704.04861
"""


import torch.nn as nn
import torch
# 基本卷积类
class BasicConv2dBlock(nn.Module):
    # 构造方法
    def __init__(self,input_channels, output_channels, kernel_size,downsample = True, **kwargs):
        """基本卷积模块

        Args:
            input_channels        (int): 输入通道数
            output_channels       (int): 输出通道数
            kernel_size           (int): 卷积核大小
            downsample (bool, optional): 是否进行下采样(一些比较小的图片如果下采样后，到后面结构太小了). Defaults to True.
        """
        super(BasicConv2dBlock,self).__init__()
        # 判断是否进行下采样
        stride = 2 if downsample else 1
        # 卷积
        self.conv   = nn.Conv2d(input_channels,output_channels,kernel_size,stride=stride,**kwargs)
        # 批量归一化
        self.bn     = nn.BatchNorm2d(output_channels)
        # 激活函数
        self.relu   = nn.ReLU(inplace=True)
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x
# 深度可分离卷积
class DepthSeperabelConv2dBlock(nn.Module):
    def __init__(self,input_channels, output_channels, kernel_size, **kwargs):
        super(DepthSeperabelConv2dBlock,self).__init__()
        # 深度卷积
        self.depth_wise = nn.Sequential(
            nn.Conv2d(input_channels,input_channels,kernel_size,groups=input_channels,**kwargs),
            nn.BatchNorm2d(input_channels),
            nn.ReLU(inplace=True)
        )
        # 逐点卷积
        self.point_wise = nn.Sequential(
            nn.Conv2d(input_channels,output_channels,1),
            nn.BatchNorm2d(output_channels),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        x = self.depth_wise(x)
        x = self.point_wise(x)

        return x
class MobileNet(nn.Module):
    """
    Args:
        width multipler: The role of the width multiplier α is to thin
                        a network uniformly at each layer. For a given
                        layer and width multiplier α, the number of
                        input channels M becomes αM and the number of
                        output channels N becomes αN.
    """

    def __init__(self, width_multiplier=1, class_num=100):
        super(MobileNet,self).__init__()

        alpha = width_multiplier

        self.stem = nn.Sequential(
            BasicConv2dBlock(3,int(32 * alpha),kernel_size = 3,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(32 * alpha),int(64 * alpha),kernel_size = 3,padding = 1,bias = False)
        )

        self.conv1 = nn.Sequential(
            DepthSeperabelConv2dBlock(int(64 * alpha),int(128 * alpha),kernel_size = 3,stride = 2,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(128 * alpha),int(128 * alpha),kernel_size = 3,padding = 1,bias = False)
        )

        self.conv2 = nn.Sequential(
            DepthSeperabelConv2dBlock(int(128 * alpha),int(256 * alpha),kernel_size = 3,stride = 2,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(256 * alpha),int(256 * alpha),kernel_size = 3,padding = 1,bias = False)
        )

        self.conv3 = nn.Sequential(
            DepthSeperabelConv2dBlock(int(256 * alpha),int(512 * alpha),kernel_size = 3,stride = 2,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(512 * alpha),int(512 * alpha),kernel_size = 3,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(512 * alpha),int(512 * alpha),kernel_size = 3,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(512 * alpha),int(512 * alpha),kernel_size = 3,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(512 * alpha),int(512 * alpha),kernel_size = 3,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(512 * alpha),int(512 * alpha),kernel_size = 3,padding = 1,bias = False)
        )
        self.conv4 = nn.Sequential(
            DepthSeperabelConv2dBlock(int(512 * alpha),int(1024 * alpha),kernel_size = 3,stride = 2,padding = 1,bias = False),
            DepthSeperabelConv2dBlock(int(1024 * alpha),int(1024 * alpha),kernel_size = 3,padding = 1,bias = False)
        )
        self.fc = nn.Linear(int(1024 * alpha), class_num)
        self.avg = nn.AdaptiveAvgPool2d(1)
    def forward(self, x):
        x = self.stem(x)

        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)

        x = self.avg(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

def mobilenet(alpha=1, class_num=100):
    return MobileNet(alpha, class_num)