text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import os
  import math
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.utils.model_zoo as model_zoo
  from torch.nn import init
  
  
  __all__ = ["MobileNetV3", "mobilenet_v3_large", "mobilenet_v3_small"]
  
  model_urls = {
      "mobilenet_v3_large": "https://download.pytorch.org/models/mobilenet_v3_large-8738ca79.pth",
      "mobilenet_v3_small": "https://download.pytorch.org/models/mobilenet_v3_small-047dcff4.pth",
  }
  
  class hswish(nn.Module):
      def forward(self, x):
          out = x * F.relu6(x + 3, inplace=True) / 6
          return out
  
  
  class hsigmoid(nn.Module):
      def forward(self, x):
          out = F.relu6(x + 3, inplace=True) / 6
          return out
  
  
  class SeModule(nn.Module):
      def __init__(self, in_size, reduction=4):
          super(SeModule, self).__init__()
          self.se = nn.Sequential(
              nn.AdaptiveAvgPool2d(1),
              nn.Conv2d(in_size, in_size // reduction, kernel_size=1, stride=1, padding=0, bias=False),
              nn.BatchNorm2d(in_size // reduction),
              nn.ReLU(inplace=True),
              nn.Conv2d(in_size // reduction, in_size, kernel_size=1, stride=1, padding=0, bias=False),
              nn.BatchNorm2d(in_size),
              hsigmoid()
          )
  
      def forward(self, x):
          return x * self.se(x)
  
  def fill_up_weights(up):
      w = up.weight.data
      f = math.ceil(w.size(2) / 2)
      c = (2 * f - 1 - f % 2) / (2. * f)
      for i in range(w.size(2)):
          for j in range(w.size(3)):
              w[0, 0, i, j] = \
                  (1 - math.fabs(i / f - c)) * (1 - math.fabs(j / f - c))
      for c in range(1, w.size(0)):
          w[c, 0, :, :] = w[0, 0, :, :] 
  
  def fill_fc_weights(layers):
      for m in layers.modules():
          if isinstance(m, nn.Conv2d):
              nn.init.normal_(m.weight, std=0.001)
              # torch.nn.init.kaiming_normal_(m.weight.data, nonlinearity='relu')
              # torch.nn.init.xavier_normal_(m.weight.data)
              if m.bias is not None:
                  nn.init.constant_(m.bias, 0)
  
  class Block(nn.Module):
      '''expand + depthwise + pointwise'''
      def __init__(self, kernel_size, in_size, expand_size, out_size, nolinear, semodule, stride):
          super(Block, self).__init__()
          self.stride = stride
          self.se = semodule
          self.conv1 = nn.Conv2d(in_size, expand_size, kernel_size=1, stride=1, padding=0, bias=False)
          self.bn1 = nn.BatchNorm2d(expand_size)
          self.nolinear1 = nolinear
          self.conv2 = nn.Conv2d(expand_size, expand_size, kernel_size=kernel_size, stride=stride,
                                 padding=kernel_size // 2, groups=expand_size, bias=False)
          self.bn2 = nn.BatchNorm2d(expand_size)
          self.nolinear2 = nolinear
          self.conv3 = nn.Conv2d(expand_size, out_size, kernel_size=1, stride=1, padding=0, bias=False)
          self.bn3 = nn.BatchNorm2d(out_size)
  
          self.shortcut = nn.Sequential()
          if stride == 1 and in_size != out_size:
              self.shortcut = nn.Sequential(
                  nn.Conv2d(in_size, out_size, kernel_size=1, stride=1, padding=0, bias=False),
                  nn.BatchNorm2d(out_size),
              )
  
      def forward(self, x):
          out = self.nolinear1(self.bn1(self.conv1(x)))
          out = self.nolinear2(self.bn2(self.conv2(out)))
          out = self.bn3(self.conv3(out))
          if self.se != None:
              out = self.se(out)
          out = out + self.shortcut(x) if self.stride == 1 else out
          return out
  
  
  class MobileNetV3_Large(nn.Module):
      def __init__(self, heads, head_conv, final_kernel):
          
          super(MobileNetV3_Large, self).__init__()
          self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=False)
          self.bn1 = nn.BatchNorm2d(16)
          self.hs1 = hswish()
  
          self.bneck = nn.Sequential(
              Block(3, 16, 16, 16, nn.ReLU(inplace=True), None, 1),
              Block(3, 16, 64, 24, nn.ReLU(inplace=True), None, 2),
              Block(3, 24, 72, 24, nn.ReLU(inplace=True), None, 1),
          )
          self.bneck1 = nn.Sequential(
              Block(5, 24, 72, 40, nn.ReLU(inplace=True), SeModule(40), 2),
              Block(5, 40, 120, 40, nn.ReLU(inplace=True), SeModule(40), 1),
              Block(5, 40, 120, 40, nn.ReLU(inplace=True), SeModule(40), 1),
          )
          self.bneck2 = nn.Sequential(
              Block(3, 40, 240, 80, hswish(), None, 2),
              Block(3, 80, 200, 80, hswish(), None, 1),
              Block(3, 80, 184, 80, hswish(), None, 1),
              Block(3, 80, 184, 80, hswish(), None, 1),
              Block(3, 80, 480, 112, hswish(), SeModule(112), 1),
              Block(3, 112, 672, 112, hswish(), SeModule(112), 1),
              Block(5, 112, 672, 160, hswish(), SeModule(160), 1),
          )
          self.bneck3 = nn.Sequential(
              Block(5, 160, 672, 160, hswish(), SeModule(160), 2),
              Block(5, 160, 960, 160, hswish(), SeModule(160), 1),
          )
  
          self.conv2 = nn.Conv2d(160, 960, kernel_size=1, stride=1, padding=0, bias=False)
          self.bn2 = nn.BatchNorm2d(960)
          self.hs2 = hswish()
  
          self.init_params()
  
          # ??
          self.inplanes = 64
          self.deconv3 = nn.ConvTranspose2d(960, self.inplanes, 8*2, stride=8, padding= 8//2, output_padding=0, 
                                          groups=self.inplanes, bias=False)
          self.bn3 = nn.BatchNorm2d(64)
          self.hs3 = hswish()
  
          self.heads = heads
          for head in self.heads:
              classes = self.heads[head]
              if head_conv > 0:
                  fc = nn.Sequential(
                    nn.Conv2d(64, head_conv,
                      kernel_size=3, padding=1, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(head_conv, classes, 
                      kernel_size=1, stride=1, 
                      padding=0, bias=True))
                  if 'hm' in head:
                      fc[-1].bias.data.fill_(-2.19)
                  else:
                      fill_fc_weights(fc)
              else:
                  fc = nn.Conv2d(64, classes, 
                    kernel_size=1, stride=1, 
                    padding=0, bias=True)
                  if 'hm' in head:
                      fc.bias.data.fill_(-2.19)
                  else:
                      fill_fc_weights(fc)
              self.__setattr__(head, fc)
  
      def init_params(self):
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  init.kaiming_normal_(m.weight, mode='fan_out')
                  if m.bias is not None:
                      init.constant_(m.bias, 0)
              elif isinstance(m, nn.BatchNorm2d):
                  init.constant_(m.weight, 1)
                  init.constant_(m.bias, 0)
              elif isinstance(m, nn.Linear):
                  init.normal_(m.weight, std=0.001)
                  if m.bias is not None:
                      init.constant_(m.bias, 0)
          
      def forward(self, x):
          out = self.hs1(self.bn1(self.conv1(x)))
          out = self.bneck(out)
          out = self.bneck1(out)
          out = self.bneck2(out)
          out = self.bneck3(out)
          out = self.hs2(self.bn2(self.conv2(out)))
          out = self.hs3(self.bn3(self.deconv3(out)))
          
          ret = {}
          for head in self.heads:
              ret[head] = self.__getattr__(head)(out)
          return [ret]
  
  class MobileNetV3_Small(nn.Module):
      def __init__(self, heads, head_conv, final_kernel):
          super(MobileNetV3_Small, self).__init__()
          self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=False)
          self.bn1 = nn.BatchNorm2d(16)
          self.hs1 = hswish()
  
          self.bneck = nn.Sequential(
              Block(3, 16, 16, 16, nn.ReLU(inplace=True), SeModule(16), 2),
              Block(3, 16, 72, 24, nn.ReLU(inplace=True), None, 2),
              Block(3, 24, 88, 24, nn.ReLU(inplace=True), None, 1),
              Block(5, 24, 96, 40, hswish(), SeModule(40), 2),
              Block(5, 40, 240, 40, hswish(), SeModule(40), 1),
              Block(5, 40, 240, 40, hswish(), SeModule(40), 1),
              Block(5, 40, 120, 48, hswish(), SeModule(48), 1),
              Block(5, 48, 144, 48, hswish(), SeModule(48), 1),
              Block(5, 48, 288, 96, hswish(), SeModule(96), 2),
              Block(5, 96, 576, 96, hswish(), SeModule(96), 1),
              Block(5, 96, 576, 96, hswish(), SeModule(96), 1),
          )
  
  
          self.conv2 = nn.Conv2d(96, 576, kernel_size=1, stride=1, padding=0, bias=False)
          self.bn2 = nn.BatchNorm2d(576) # ??
          self.hs2 = hswish()
  
          self.init_params()
  
          self.inplanes = 64
          self.deconv3 = nn.ConvTranspose2d(576, self.inplanes, 8*2, stride=8, padding= 8//2, output_padding=0, 
                                          groups=self.inplanes, bias=False)
          self.bn3 = nn.BatchNorm2d(64)
          self.hs3 = hswish()
  
          self.heads = heads
          for head in self.heads:
              classes = self.heads[head]
              if head_conv > 0:
                  fc = nn.Sequential(
                    nn.Conv2d(64, head_conv,
                      kernel_size=3, padding=1, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(head_conv, classes, 
                      kernel_size=1, stride=1, 
                      padding=0, bias=True))
                  if 'hm' in head:
                      fc[-1].bias.data.fill_(-2.19)
                  else:
                      fill_fc_weights(fc)
              else:
                  fc = nn.Conv2d(64, classes, 
                    kernel_size=1, stride=1, 
                    padding=0, bias=True)
                  if 'hm' in head:
                      fc.bias.data.fill_(-2.19)
                  else:
                      fill_fc_weights(fc)
              self.__setattr__(head, fc)
              
      def init_params(self):
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  init.kaiming_normal_(m.weight, mode='fan_out')
                  if m.bias is not None:
                      init.constant_(m.bias, 0)
              elif isinstance(m, nn.BatchNorm2d):
                  init.constant_(m.weight, 1)
                  init.constant_(m.bias, 0)
              elif isinstance(m, nn.Linear):
                  init.normal_(m.weight, std=0.001)
                  if m.bias is not None:
                      init.constant_(m.bias, 0)
      
  
      def forward(self, x):
          out = self.hs1(self.bn1(self.conv1(x)))
          out = self.bneck(out)
          out = self.hs2(self.bn2(self.conv2(out)))
          out = self.hs3(self.bn3(self.deconv3(out)))
          ret = {}
          for head in self.heads:
              ret[head] = self.__getattr__(head)(out)
          return [ret]
  
  mobilenetv3_spec = {'large': MobileNetV3_Large, 
                      'small': MobileNetV3_Small}
  
  def get_mobilenet_v3(num_layers, heads, head_conv):
      select_model = mobilenetv3_spec[num_layers]
      model = select_model(heads=heads, head_conv=head_conv, final_kernel=1)
      return model
  
  