FB

nnet/xvector.py

@@ -384,6 +384,76 @@ class MfccFrontEnd(torch.nn.Module):
                 mfcc = self.CMVN(mfcc)
         return mfcc
 
+
+class MelSpecFrontEnd(torch.nn.Module):
+    """
+
+    """
+
+    def __init__(self,
+                 pre_emphasis=0.97,
+                 sample_rate=16000,
+                 n_fft=1024,
+                 f_min=90,
+                 f_max=7600,
+                 win_length=1024,
+                 window_fn=torch.hann_window,
+                 hop_length=256,
+                 power=2.0,
+                 n_mels=80):
+
+        super(MelSpecFrontEnd, self).__init__()
+
+        self.pre_emphasis = pre_emphasis
+        self.sample_rate = sample_rate
+        self.n_fft = n_fft
+        self.f_min = f_min
+        self.f_max = f_max
+        self.win_length = win_length
+        self.window_fn=window_fn
+        self.hop_length = hop_length
+        self.power=power
+        self.window_fn = window_fn
+        self.n_mels = n_mels
+
+        self.PreEmphasis = PreEmphasis(self.pre_emphasis)
+
+        self.melkwargs = {"n_fft":self.n_fft,
+                          "f_min":self.f_min,
+                          "f_max":self.f_max,
+                          "win_length":self.win_length,
+                          "window_fn":self.window_fn,
+                          "hop_length":self.hop_length,
+                          "power":self.power,
+                          "n_mels":self.n_mels}
+
+        self.MelSpec = torchaudio.transforms.MelSpectrogram(sample_rate=self.sample_rate,
+                                                            n_fft=self.melkwargs['n_fft'],
+                                                            f_min=self.melkwargs['f_min'],
+                                                            f_max=self.melkwargs['f_max'],
+                                                            win_length=self.melkwargs['win_length'],
+                                                            hop_length=self.melkwargs['hop_length'],
+                                                            window_fn=self.melkwargs['window_fn'],
+                                                            power=self.melkwargs['power'],
+                                                            n_mels=self.melkwargs['n_mels'])
+
+        self.CMVN = torch.nn.InstanceNorm1d(self.n_mels)
+
+    def forward(self, x):
+        """
+
+        :param x:
+        :return:
+        """
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(enabled=False):
+                out = self.PreEmphasis(x)
+                out = self.MelSpec(out)+1e-6
+                out = torch.log(out)
+                out = self.CMVN(out)
+        return out
+
+
 class Xtractor(torch.nn.Module):
     """
     Class that defines an x-vector extractor based on 5 convolutional layers and a mean standard deviation pooling
@@ -497,7 +567,7 @@ class Xtractor(torch.nn.Module):
 
         elif model_archi == "fastresnet34":
 
-            self.preprocessor = MfccFrontEnd()
+            self.preprocessor = MelSpecFrontEnd()
             self.sequence_network = PreFastResNet34()
             self.embedding_size = 256