Update modeling_wav2vec2_spkreg.py

modeling_wav2vec2_spkreg.py

@@ -578,7 +578,7 @@ class AAMSoftmaxLoss(nn.Module):
     def __init__(
         self, 
         scale: float = 30.0, 
-        margin: float = 0.35, 
+        margin: float = 0.2, 
         easy_margin: bool = False, 
         label_smoothing: float = 0.0, 
         reduction: str = "mean"
@@ -596,9 +596,6 @@ class AAMSoftmaxLoss(nn.Module):
         self.easy_margin = easy_margin
         self.label_smoothing = label_smoothing
         self.reduction = reduction
-
-        self.cos_m = math.cos(self.margin)
-        self.sin_m = math.sin(self.margin)
         
     def forward(
         self, 
@@ -614,11 +611,23 @@ class AAMSoftmaxLoss(nn.Module):
         """
         _, num_labels = inputs.shape
         # `inputs` are the outputs from AngularLinear()
-        cos_theta = torch.clamp(inputs, -1.0, 1.0)
-        sin_theta = torch.sqrt(1.0 - torch.pow(cos_theta, 2))
-        psi = cos_theta * self.cos_m - sin_theta * self.sin_m # cos(theta + m)
+        epsilon = 1e-6
         # theta = torch.acos(cos_theta)
         # psi = torch.cos(theta + self.margin)
+        cos_theta = torch.clamp(inputs, -1.0 + epsilon, 1.0 - epsilon)
+        sin_theta = torch.sqrt(1.0 - torch.pow(cos_theta, 2))
+        sin_theta = torch.clamp(sin_theta, 0.0 + epsilon, 1.0 - epsilon)
+
+        cos_m = math.cos(self.margin)
+        sin_m = math.sin(self.margin)
+        psi = cos_theta * cos_m - sin_theta * sin_m # cos(theta + m)
+
+        if self.easy_margin:
+            psi = torch.where(cos_theta > 0, psi, cos_theta)
+        else:
+            # Make the function cos(theta+m) monotonic decreasing while theta in [0°, 180°]
+            psi = torch.where((cos_theta - math.cos(math.pi - self.margin)) > 0, psi, cos_theta - self.margin)
+
         one_hot = nn.functional.one_hot(targets, num_labels)
         outputs = self.scale * torch.where(one_hot.bool(), psi, cos_theta)
         loss = F.cross_entropy(