mobilenet v1结构介绍

mobilenet v1

mobilenetv1中的深度可分离卷积模块

mobilenetv1完整结构：

计算量

参考：MobileNetv1论文详解.md

代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from torch import flatten

class MobilnetV1Block(nn.Module):
    """Depthwise conv + Pointwise conv"""

    def __init__(self, in_channels, out_channels, stride=1):
        super(MobilnetV1Block, self).__init__()
        # dw conv kernel shape is (in_channels, 1, ksize, ksize)
        self.dw = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3,
                      stride=stride, padding=1, groups=4, bias=False),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True)
        )
        # print(self.dw[0].weight.shape)  # print dw conv kernel shape
        self.pw = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=1,
                      stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.dw(x)
        x = self.pw(x)
        return x


def convbn_relu(in_channels, out_channels, stride=2):
    return nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride,
                                   padding=1, bias=False),
                         nn.BatchNorm2d(out_channels),
                         nn.ReLU(inplace=True))


class MobileNetV1(nn.Module):
    # (32, 64, 1) means MobilnetV1Block in_channnels is 32, out_channels is 64, no change in map size.
    stage_cfg = [(32, 64, 1), 
           (64, 128, 2), (128, 128, 1),     # stage1 conv
           (128, 256, 2), (256, 256, 1),    # stage2 conv
           (256, 512, 2), (512, 512, 1), (512, 512, 1), (512, 512, 1), (512, 512, 1), (512, 512, 1), # stage3 conv
           (512, 1024, 2), (1024, 1024, 1)  # stage4 conv
    ]
    def __init__(self, num_classes=1000):
        super(MobileNetV1, self).__init__()
        self.first_conv = convbn_relu(3, 32, 2)  # Input image size reduced by half
        self.stage_layers = self._make_layers(in_channels=32)
        self.linear = nn.Linear(1024, num_classes)  # 全连接层

    def _make_layers(self, in_channels):
        layers = []
        for x in self.stage_cfg:
            in_channels = x[0]
            out_channels = x[1]
            stride = x[2]
            layers.append(MobilnetV1Block(in_channels, out_channels, stride))
            in_channels = out_channels
        return nn.Sequential(*layers)

    def forward(self, x):
        """Feature map shape(h、w) is 224 -> 112 -> 56 -> 28 -> 14 -> 7 -> 1"""
        x = self.first_conv(x)
        x = self.stage_layers(x)

        x = F.avg_pool2d(x, 7)  # x shape is 7*7
        x = flatten(x, 1)       # x = x.view(x.size(0), -1)
        x = self.linear(x)

        return x


if __name__ == "__main__":
    model = MobileNetV1()
    model.eval()                  # set the model to inference mode
    input_data = torch.rand(1, 3, 224, 224)
    outputs = model(input_data)
    print("Model output size is", outputs.size())