model.py

import torch.nn as nn
import math


class Bottleneck(nn.Module):

    def __init__(self, inplanes, outplanes, stride=1, expansion=6):
        super(Bottleneck, self).__init__()
        self.stride = stride
        self.conv1 = nn.Conv2d(inplanes, inplanes*expansion, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(inplanes*expansion)
        self.conv2 = nn.Conv2d(inplanes*expansion, inplanes*expansion, kernel_size=3, stride=stride,
                               padding=1, bias=False, groups=inplanes*expansion)
        self.bn2 = nn.BatchNorm2d(inplanes*expansion)
        self.conv3 = nn.Conv2d(inplanes*expansion, outplanes, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(outplanes)
        self.relu = nn.ReLU6(inplace=True)
        self.use_residual = self.stride == 1 and inplanes == outplanes

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.use_residual:
            return residual+out
        return out


class MobileNetV2(nn.Module):

    def __init__(self, block, layers, input_size=224, width_multi=1.,expansion=6, num_classes=1000):
        self.inplanes = int(32*width_multi)
        super(MobileNetV2, self).__init__()
        self.conv1 = nn.Conv2d(3, int(32*width_multi), kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(int(32*width_multi))
        self.relu = nn.ReLU6(inplace=True)
        self.layer1 = self._make_layer(block, int(16*width_multi), layers[0], stride=1, expansion=1)
        self.layer2 = self._make_layer(block, int(24*width_multi), layers[1], stride=2, expansion=expansion)
        self.layer3 = self._make_layer(block, int(32*width_multi), layers[2], stride=2, expansion=expansion)
        self.layer4 = self._make_layer(block, int(64*width_multi), layers[3], stride=2, expansion=expansion)
        self.layer5 = self._make_layer(block, int(96*width_multi), layers[4], stride=1, expansion=expansion)
        self.layer6 = self._make_layer(block, int(160*width_multi), layers[5], stride=2, expansion=expansion)
        self.layer7 = self._make_layer(block, int(320*width_multi), layers[6], stride=1, expansion=expansion)
        last_channel = int(1280*width_multi) if width_multi>1. else 1280
        self.conv8 = nn.Conv2d(int(320*width_multi), last_channel, kernel_size=1, stride=1, bias=False)
        self.avgpool = nn.AvgPool2d(int(input_size/32), stride=1)
        self.conv9 = nn.Conv2d(last_channel, num_classes, kernel_size=1, stride=1, bias=False)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0]*m.kernel_size[1]*m.out_channels
                m.weight.data.normal_(0, math.sqrt(2./n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride, expansion):
        layers = []
        layers.append(block(self.inplanes, planes, stride=stride, expansion=expansion))
        self.inplanes = planes
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, expansion=expansion))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = self.layer6(x)
        x = self.layer7(x)

        x = self.conv8(x)
        x = self.avgpool(x)
        x = self.conv9(x)
        x = x.view(x.size(0), -1)

        return x


def mobilenetv2_19(**kwargs):
    model = MobileNetV2(Bottleneck, [1, 2, 3, 4, 3, 3, 1], **kwargs)

    return model