htrnguyen/golf-tech-analysis

1import torch.nn as nn2import math3 4"""5https://github.com/tonylins/pytorch-mobilenet-v26"""7 8def conv_bn(inp, oup, stride):9    return nn.Sequential(10        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),11        nn.BatchNorm2d(oup),12        nn.ReLU6(inplace=True)13    )14 15 16def conv_1x1_bn(inp, oup):17    return nn.Sequential(18        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),19        nn.BatchNorm2d(oup),20        nn.ReLU6(inplace=True)21    )22 23 24class InvertedResidual(nn.Module):25    def __init__(self, inp, oup, stride, expand_ratio):26        super(InvertedResidual, self).__init__()27        self.stride = stride28        assert stride in [1, 2]29 30        hidden_dim = round(inp * expand_ratio)31        self.use_res_connect = self.stride == 1 and inp == oup32 33        if expand_ratio == 1:34            self.conv = nn.Sequential(35                # dw36                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),37                nn.BatchNorm2d(hidden_dim),38                nn.ReLU6(inplace=True),39                # pw-linear40                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),41                nn.BatchNorm2d(oup),42            )43        else:44            self.conv = nn.Sequential(45                # pw46                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),47                nn.BatchNorm2d(hidden_dim),48                nn.ReLU6(inplace=True),49                # dw50                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),51                nn.BatchNorm2d(hidden_dim),52                nn.ReLU6(inplace=True),53                # pw-linear54                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),55                nn.BatchNorm2d(oup),56            )57 58    def forward(self, x):59        if self.use_res_connect:60            return x + self.conv(x)61        else:62            return self.conv(x)63 64 65class MobileNetV2(nn.Module):66    def __init__(self, n_class=1000, input_size=224, width_mult=1.):67        super(MobileNetV2, self).__init__()68        block = InvertedResidual69        min_depth = 1670        input_channel = 3271        last_channel = 128072        interverted_residual_setting = [73            # t, c, n, s74            [1, 16, 1, 1],75            [6, 24, 2, 2],76            [6, 32, 3, 2],77            [6, 64, 4, 2],78            [6, 96, 3, 1],79            [6, 160, 3, 2],80            [6, 320, 1, 1],81        ]82 83        # building first layer84        assert input_size % 32 == 085        input_channel = int(input_channel * width_mult) if width_mult >= 1.0 else input_channel86        self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel87        self.features = [conv_bn(3, input_channel, 2)]88        # building inverted residual blocks89        for t, c, n, s in interverted_residual_setting:90            output_channel = max(int(c * width_mult), min_depth)91            for i in range(n):92                if i == 0:93                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t))94                else:95                    self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))96                input_channel = output_channel97        # building last several layers98        self.features.append(conv_1x1_bn(input_channel, self.last_channel))99        # make it nn.Sequential100        self.features = nn.Sequential(*self.features)101 102        # building classifier103        self.classifier = nn.Sequential(104            nn.Dropout(0.2),105            nn.Linear(self.last_channel, n_class),106        )107 108        self._initialize_weights()109 110    def forward(self, x):111        x = self.features(x)112        x = x.mean(3).mean(2)113        x = self.classifier(x)114        return x115 116    def _initialize_weights(self):117        for m in self.modules():118            if isinstance(m, nn.Conv2d):119                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels120                m.weight.data.normal_(0, math.sqrt(2. / n))121                if m.bias is not None:122                    m.bias.data.zero_()123            elif isinstance(m, nn.BatchNorm2d):124                m.weight.data.fill_(1)125                m.bias.data.zero_()126            elif isinstance(m, nn.Linear):127                n = m.weight.size(1)128                m.weight.data.normal_(0, 0.01)129                m.bias.data.zero_()130