linter and test updates for speaker_encoder, gmm_Attention

speaker_encoder/loss.py

@@ -5,9 +5,8 @@ import torch.nn.functional as F
 
 # adapted from https://github.com/cvqluu/GE2E-Loss
 class GE2ELoss(nn.Module):
-
-    def __init__(self, init_w=10.0, init_b=-5.0, loss_method='softmax'):
-        '''
+    def __init__(self, init_w=10.0, init_b=-5.0, loss_method="softmax"):
+        """
         Implementation of the Generalized End-to-End loss defined in https://arxiv.org/abs/1710.10467 [1]
         Accepts an input of size (N, M, D)
             where N is the number of speakers in the batch,
@@ -16,24 +15,27 @@ class GE2ELoss(nn.Module):
         Args:
             - init_w (float): defines the initial value of w in Equation (5) of [1]
             - init_b (float): definies the initial value of b in Equation (5) of [1]
-        '''
+        """
         super(GE2ELoss, self).__init__()
+        # pylint: disable=E1102
         self.w = nn.Parameter(torch.tensor(init_w))
+        # pylint: disable=E1102
         self.b = nn.Parameter(torch.tensor(init_b))
         self.loss_method = loss_method
 
-        assert self.loss_method in ['softmax', 'contrast']
+        assert self.loss_method in ["softmax", "contrast"]
 
-        if self.loss_method == 'softmax':
+        if self.loss_method == "softmax":
             self.embed_loss = self.embed_loss_softmax
-        if self.loss_method == 'contrast':
+        if self.loss_method == "contrast":
             self.embed_loss = self.embed_loss_contrast
 
+    # pylint: disable=R0201
     def calc_new_centroids(self, dvecs, centroids, spkr, utt):
-        '''
+        """
         Calculates the new centroids excluding the reference utterance
-        '''
-        excl = torch.cat((dvecs[spkr, :utt], dvecs[spkr, utt+1:]))
+        """
+        excl = torch.cat((dvecs[spkr, :utt], dvecs[spkr, utt + 1 :]))
         excl = torch.mean(excl, 0)
         new_centroids = []
         for i, centroid in enumerate(centroids):
@@ -44,26 +46,36 @@ class GE2ELoss(nn.Module):
         return torch.stack(new_centroids)
 
     def calc_cosine_sim(self, dvecs, centroids):
-        '''
+        """
         Make the cosine similarity matrix with dims (N,M,N)
-        '''
+        """
         cos_sim_matrix = []
         for spkr_idx, speaker in enumerate(dvecs):
             cs_row = []
             for utt_idx, utterance in enumerate(speaker):
                 new_centroids = self.calc_new_centroids(
-                    dvecs, centroids, spkr_idx, utt_idx)
+                    dvecs, centroids, spkr_idx, utt_idx
+                )
                 # vector based cosine similarity for speed
-                cs_row.append(torch.clamp(torch.mm(utterance.unsqueeze(1).transpose(0, 1), new_centroids.transpose(
-                    0, 1)) / (torch.norm(utterance) * torch.norm(new_centroids, dim=1)), 1e-6))
+                cs_row.append(
+                    torch.clamp(
+                        torch.mm(
+                            utterance.unsqueeze(1).transpose(0, 1),
+                            new_centroids.transpose(0, 1),
+                        )
+                        / (torch.norm(utterance) * torch.norm(new_centroids, dim=1)),
+                        1e-6,
+                    )
+                )
             cs_row = torch.cat(cs_row, dim=0)
             cos_sim_matrix.append(cs_row)
         return torch.stack(cos_sim_matrix)
 
+    # pylint: disable=R0201
     def embed_loss_softmax(self, dvecs, cos_sim_matrix):
-        '''
+        """
         Calculates the loss on each embedding $L(e_{ji})$ by taking softmax
-        '''
+        """
         N, M, _ = dvecs.shape
         L = []
         for j in range(N):
@@ -74,10 +86,11 @@ class GE2ELoss(nn.Module):
             L.append(L_row)
         return torch.stack(L)
 
+    # pylint: disable=R0201
     def embed_loss_contrast(self, dvecs, cos_sim_matrix):
-        ''' 
+        """
         Calculates the loss on each embedding $L(e_{ji})$ by contrast loss with closest centroid
-        '''
+        """
         N, M, _ = dvecs.shape
         L = []
         for j in range(N):
@@ -85,17 +98,21 @@ class GE2ELoss(nn.Module):
             for i in range(M):
                 centroids_sigmoids = torch.sigmoid(cos_sim_matrix[j, i])
                 excl_centroids_sigmoids = torch.cat(
-                    (centroids_sigmoids[:j], centroids_sigmoids[j+1:]))
+                    (centroids_sigmoids[:j], centroids_sigmoids[j + 1 :])
+                )
                 L_row.append(
-                    1. - torch.sigmoid(cos_sim_matrix[j, i, j]) + torch.max(excl_centroids_sigmoids))
+                    1.0
+                    - torch.sigmoid(cos_sim_matrix[j, i, j])
+                    + torch.max(excl_centroids_sigmoids)
+                )
             L_row = torch.stack(L_row)
             L.append(L_row)
         return torch.stack(L)
 
     def forward(self, dvecs):
-        '''
+        """
         Calculates the GE2E loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats)
-        '''
+        """
         centroids = torch.mean(dvecs, 1)
         cos_sim_matrix = self.calc_cosine_sim(dvecs, centroids)
         torch.clamp(self.w, 1e-6)