MobileNet

In [ ]:

import torch.nn as nn

import torch.nn as nn

In [ ]:

class MobileBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super(MobileBlock, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False),
            nn.BatchNorm2d(in_channels),
            nn.ReLU6(),
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU6()
        )
    def forward(self, x):
        return self.net(x)

class MobileBlock(nn.Module): def __init__(self, in_channels, out_channels, stride=1): super(MobileBlock, self).__init__() self.net = nn.Sequential( nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False), nn.BatchNorm2d(in_channels), nn.ReLU6(), nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU6() ) def forward(self, x): return self.net(x)

In [ ]:

class MobileNet(nn.Module):
    def __init__(self, num_classes=10):
        super(MobileNet, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU6(),
            MobileBlock(32, 64, 1),
            MobileBlock(64, 128, 2),
            MobileBlock(128, 128, 1),
            MobileBlock(128, 256, 2),
            MobileBlock(256, 256, 1),
            MobileBlock(256, 512, 2),
            MobileBlock(512, 512, 1),
            MobileBlock(512, 512, 1),
            MobileBlock(512, 512, 1),
            MobileBlock(512, 512, 1),
            MobileBlock(512, 512, 1),
            MobileBlock(512, 1024, 2),
            MobileBlock(1024, 1024, 1),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(in_features=1024, out_features=num_classes)
        )

    def forward(self, x):
        return self.net(x)

class MobileNet(nn.Module): def __init__(self, num_classes=10): super(MobileNet, self).__init__() self.net = nn.Sequential( nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU6(), MobileBlock(32, 64, 1), MobileBlock(64, 128, 2), MobileBlock(128, 128, 1), MobileBlock(128, 256, 2), MobileBlock(256, 256, 1), MobileBlock(256, 512, 2), MobileBlock(512, 512, 1), MobileBlock(512, 512, 1), MobileBlock(512, 512, 1), MobileBlock(512, 512, 1), MobileBlock(512, 512, 1), MobileBlock(512, 1024, 2), MobileBlock(1024, 1024, 1), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(in_features=1024, out_features=num_classes) ) def forward(self, x): return self.net(x)