text: """senet in pytorch
  
  
  
  [1] Jie Hu, Li Shen, Samuel Albanie, Gang Sun, Enhua Wu
  
      Squeeze-and-Excitation Networks
      https://arxiv.org/abs/1709.01507
  """
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  class SepConv(nn.Module):
      def __init__(self, channel_in, channel_out, kernel_size=3, stride=2, padding=1, affine=True):
          super(SepConv, self).__init__()
          self.op = nn.Sequential(
              nn.Conv2d(channel_in, channel_in, kernel_size=kernel_size, stride=stride, padding=padding, groups=channel_in, bias=False),
              nn.Conv2d(channel_in, channel_in, kernel_size=1, padding=0, bias=False),
              nn.BatchNorm2d(channel_in, affine=affine),
              nn.ReLU(inplace=False),
              nn.Conv2d(channel_in, channel_in, kernel_size=kernel_size, stride=1, padding=padding, groups=channel_in, bias=False),
              nn.Conv2d(channel_in, channel_out, kernel_size=1, padding=0, bias=False),
              nn.BatchNorm2d(channel_out, affine=affine),
              nn.ReLU(inplace=False),
          )
  
      def forward(self, x):
          return self.op(x)
  
  class MainClassifier(nn.Module):
      def __init__(self, channel, num_classes=100):
          super(MainClassifier, self).__init__()
          self.pool = nn.AdaptiveAvgPool2d(1)
          self.fc = nn.Linear(channel, num_classes)
  
      def forward(self, x):
          x = self.pool(x)
          x = x.view(x.size(0), -1)
          x = self.fc(x)
          return x
  
  
  class DeepwiseAuxiliaryClassifier(nn.Module):
  
      def __init__(self, channel, num_classes=100, downsample=0):
          super(DeepwiseAuxiliaryClassifier, self).__init__()
          self.channel = channel
          self.pool = nn.AdaptiveAvgPool2d(1)
          self.downsample = downsample
          self.layer = self._make_conv_layer()
          self.fc = nn.Linear(self.channel, num_classes)
  
  
      def _make_conv_layer(self):
          layer_list = []
          for i in range(self.downsample):
              layer_list.append(SepConv(self.channel, self.channel*2))
              self.channel *= 2
          layer_list.append(nn.AdaptiveAvgPool2d(1))
          return nn.Sequential(*layer_list)
  
      def forward(self, x):
          x = self.layer(x)
          x = x.view(x.size(0), -1)
          #x = self.fc(x)
          return x
  
  class BasicResidualSEBlock(nn.Module):
  
      expansion = 1
  
      def __init__(self, in_channels, out_channels, stride, r=16):
          super().__init__()
  
          self.residual = nn.Sequential(
              nn.Conv2d(in_channels, out_channels, 3, stride=stride, padding=1),
              nn.BatchNorm2d(out_channels),
              nn.ReLU(inplace=True),
  
              nn.Conv2d(out_channels, out_channels * self.expansion, 3, padding=1),
              nn.BatchNorm2d(out_channels * self.expansion),
              nn.ReLU(inplace=True)
          )
  
          self.shortcut = nn.Sequential()
          if stride != 1 or in_channels != out_channels * self.expansion:
              self.shortcut = nn.Sequential(
                  nn.Conv2d(in_channels, out_channels * self.expansion, 1, stride=stride),
                  nn.BatchNorm2d(out_channels * self.expansion)
              )
  
          self.squeeze = nn.AdaptiveAvgPool2d(1)
          self.excitation = nn.Sequential(
              nn.Linear(out_channels * self.expansion, out_channels * self.expansion // r),
              nn.ReLU(inplace=True),
              nn.Linear(out_channels * self.expansion // r, out_channels * self.expansion),
              nn.Sigmoid()
          )
  
      def forward(self, x):
          shortcut = self.shortcut(x)
          residual = self.residual(x)
  
          squeeze = self.squeeze(residual)
          squeeze = squeeze.view(squeeze.size(0), -1)
          excitation = self.excitation(squeeze)
          excitation = excitation.view(residual.size(0), residual.size(1), 1, 1)
  
          x = residual * excitation.expand_as(residual) + shortcut
  
          return F.relu(x)
  
  class BottleneckResidualSEBlock(nn.Module):
  
      expansion = 4
  
      def __init__(self, in_channels, out_channels, stride, r=16):
          super().__init__()
  
          self.residual = nn.Sequential(
              nn.Conv2d(in_channels, out_channels, 1),
              nn.BatchNorm2d(out_channels),
              nn.ReLU(inplace=True),
  
              nn.Conv2d(out_channels, out_channels, 3, stride=stride, padding=1),
              nn.BatchNorm2d(out_channels),
              nn.ReLU(inplace=True),
  
              nn.Conv2d(out_channels, out_channels * self.expansion, 1),
              nn.BatchNorm2d(out_channels * self.expansion),
              nn.ReLU(inplace=True)
          )
  
          self.squeeze = nn.AdaptiveAvgPool2d(1)
          self.excitation = nn.Sequential(
              nn.Linear(out_channels * self.expansion, out_channels * self.expansion // r),
              nn.ReLU(inplace=True),
              nn.Linear(out_channels * self.expansion // r, out_channels * self.expansion),
              nn.Sigmoid()
          )
  
          self.shortcut = nn.Sequential()
          if stride != 1 or in_channels != out_channels * self.expansion:
              self.shortcut = nn.Sequential(
                  nn.Conv2d(in_channels, out_channels * self.expansion, 1, stride=stride),
                  nn.BatchNorm2d(out_channels * self.expansion)
              )
  
      def forward(self, x):
  
          shortcut = self.shortcut(x)
  
          residual = self.residual(x)
          squeeze = self.squeeze(residual)
          squeeze = squeeze.view(squeeze.size(0), -1)
          excitation = self.excitation(squeeze)
          excitation = excitation.view(residual.size(0), residual.size(1), 1, 1)
  
          x = residual * excitation.expand_as(residual) + shortcut
  
          return F.relu(x)
  
  class SEResNet(nn.Module):
  
      def __init__(self, block, block_num, class_num=100):
          super().__init__()
  
          self.in_channels = 64
  
          self.pre = nn.Sequential(
              nn.Conv2d(3, 64, 3, padding=1),
              nn.BatchNorm2d(64),
              nn.ReLU(inplace=True)
          )
  
          self.stage1 = self._make_stage(block, block_num[0], 64, 1)
          self.stage2 = self._make_stage(block, block_num[1], 128, 2)
          self.stage3 = self._make_stage(block, block_num[2], 256, 2)
          self.stage4 = self._make_stage(block, block_num[3], 512, 2)
  
          self.linear = nn.Linear(self.in_channels, class_num)
          self.deepwise1 = DeepwiseAuxiliaryClassifier(channel=64 * block.expansion, downsample=3)
          self.deepwise2 = DeepwiseAuxiliaryClassifier(channel=128 * block.expansion, downsample=2)
          self.deepwise3 = DeepwiseAuxiliaryClassifier(channel=256 * block.expansion, downsample=1)
          self.deepwise4 = DeepwiseAuxiliaryClassifier(channel=512 * block.expansion, downsample=0)
  
  
      def forward(self, x):
          feature_list = []
          x = self.pre(x)
          x = self.stage1(x)
          feature_list.append(x)
          x = self.stage2(x)
          feature_list.append(x)
          x = self.stage3(x)
          feature_list.append(x)
          x = self.stage4(x)
          feature_list.append(x)
  
          x1 = self.deepwise1(feature_list[-4])
          x2 = self.deepwise2(feature_list[-3])
          x3 = self.deepwise3(feature_list[-2])
          x4 = self.deepwise4(feature_list[-1])
  
          feature = [x4, x3, x2, x1]
          x1 = self.deepwise1.fc(x1)
          x2 = self.deepwise2.fc(x2)
          x3 = self.deepwise3.fc(x3)
          x4 = self.deepwise4.fc(x4)
          return [x4, x3, x2, x1], feature
  
  
      
      def _make_stage(self, block, num, out_channels, stride):
  
          layers = []
          layers.append(block(self.in_channels, out_channels, stride))
          self.in_channels = out_channels * block.expansion
  
          while num - 1:
              layers.append(block(self.in_channels, out_channels, 1))
              num -= 1
          
          return nn.Sequential(*layers)
          
  def seresnet18():
      return SEResNet(BasicResidualSEBlock, [2, 2, 2, 2])
  
  def seresnet34():
      return SEResNet(BasicResidualSEBlock, [3, 4, 6, 3])
  
  def seresnet50():
      return SEResNet(BottleneckResidualSEBlock, [3, 4, 6, 3])
  
  def seresnet101():
      return SEResNet(BottleneckResidualSEBlock, [3, 4, 23, 3])
  
  def seresnet152():
      return SEResNet(BottleneckResidualSEBlock, [3, 8, 36, 3])