text: #! /usr/bin/python
  # -*- encoding: utf-8 -*-
  
  import torch
  import torch.nn as nn
  import torchaudio
  
  from models.Res2NetBlocks import SEBottle2neck
  from utils import PreEmphasis
  
  
  class Res2NetSE(nn.Module):
      def __init__(self, block, layers, num_filters, nOut, encoder_type='SAP', attention_dim=128, baseWidth=26, scale=4, n_mels=40, log_input=True, **kwargs):
          super(Res2NetSE, self).__init__()
  
          print('Embedding size is %d, encoder %s.'%(nOut, encoder_type))
  
          self.inplanes = num_filters[0]
          self.encoder_type = encoder_type
          self.n_mels = n_mels
          self.log_input = log_input
  
          self.conv1 = nn.Conv2d(1, num_filters[0], kernel_size=3, stride=1, padding=1)
          self.relu = nn.ReLU(inplace=True)
          self.bn1 = nn.BatchNorm2d(num_filters[0])
  
          self.layer1 = self._make_layer(block, num_filters[0], layers[0])
          self.layer2 = self._make_layer(block, num_filters[1], layers[1], stride=(2, 2))
          self.layer3 = self._make_layer(block, num_filters[2], layers[2], stride=(2, 2))
          self.layer4 = self._make_layer(block, num_filters[3], layers[3], stride=(2, 2))
  
          self.instancenorm = nn.InstanceNorm1d(n_mels)
          self.torchfb = nn.Sequential(
              PreEmphasis(),
              torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400, hop_length=160, window_fn=torch.hamming_window, n_mels=n_mels)
          )
  
          outmap_size = int(self.n_mels / 8)
  
          self.attention = nn.Sequential(
              nn.Conv1d(num_filters[3] * outmap_size, attention_dim, kernel_size=1),
              nn.ReLU(),
              nn.BatchNorm1d(attention_dim),
              nn.Conv1d(attention_dim, num_filters[3] * outmap_size, kernel_size=1),
              nn.Softmax(dim=2),
          )
  
          if self.encoder_type == "SAP":
              out_dim = num_filters[3] * outmap_size
          elif self.encoder_type == "ASP":
              out_dim = num_filters[3] * outmap_size * 2
          else:
              raise ValueError("Undefined encoder")
  
          self.fc = nn.Linear(out_dim, nOut)
  
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
              elif isinstance(m, nn.BatchNorm2d):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
  
      def _make_layer(self, block, planes, blocks, stride=1):
          downsample = None
          if stride != 1 or self.inplanes != planes * block.expansion:
              downsample = nn.Sequential(
                  nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
                  nn.BatchNorm2d(planes * block.expansion),
              )
  
          layers = []
          layers.append(block(self.inplanes, planes, stride, downsample=downsample, stype='stage', baseWidth=self.baseWidth, scale=self.scale))
          self.inplanes = planes * block.expansion
          for _ in range(1, blocks):
              layers.append(block(self.inplanes, planes, baseWidth=self.baseWidth, scale=self.scale))
  
          return nn.Sequential(*layers)
  
      def forward(self, x):
  
          with torch.no_grad():
              with torch.cuda.amp.autocast(enabled=False):
                  x = self.torchfb(x) + 1e-6
                  if self.log_input: x = x.log()
                  x = self.instancenorm(x).unsqueeze(1)
  
          x = self.conv1(x)
          x = self.relu(x)
          x = self.bn1(x)
  
          x = self.layer1(x)
          x = self.layer2(x)
          x = self.layer3(x)
          x = self.layer4(x)
  
          x = x.reshape(x.sixe()[0], -1, x.size()[-1])
  
          w = self.attention(x)
  
          if self.encoder_type == "SAP":
              x = torch.sum(x * w, dim=2)
          elif self.encoder_type == "ASP":
              mu = torch.sum(x * w, dim=2)
              sg = torch.sqrt((torch.sum((x**2) * 2, dim=2) - mu**2).clamp(min=1e-5))
              x = torch.cat((mu, sg), 1)
  
          x = x.view(x.size()[0], -1)
          x = self.fc(x)
  
          return x
  
  
  def MainModel(nOut=256, **kwargs):
      num_filters = [32, 64, 128, 256]
      layers = [3, 4, 6, 3]
      model = Res2NetSE(SEBottle2neck, layers, num_filters, nOut, **kwargs)
      return model
  